[hyperlight-guest-tracing] fixup to add tests and improve macro

- This adds tests for both the macro and the tracing crates
- Some refactoring was needed to the tracing crate to test thoroughly
- Improve macros by allowing expressions to be wrapped with trace records

Signed-off-by: Doru Blânzeanu <[email protected]>

Author: dblnz

src/hyperlight_guest/src/guest_handle/io.rs

@@ -69,17 +69,18 @@ impl GuestHandle {
         let buffer = &idb[last_element_offset_rel as usize..];
 
         // convert the buffer to T
-        hyperlight_guest_tracing_macro::trace!("Start converting buffer");
-        let type_t = match T::try_from(buffer) {
-            Ok(t) => Ok(t),
-            Err(_e) => {
-                return Err(HyperlightGuestError::new(
-                    ErrorCode::GuestError,
-                    format!("Unable to convert buffer to {}", type_name::<T>()),
-                ));
+        let type_t = hyperlight_guest_tracing_macro::trace!(
+            "convert buffer",
+            match T::try_from(buffer) {
+                Ok(t) => Ok(t),
+                Err(_e) => {
+                    return Err(HyperlightGuestError::new(
+                        ErrorCode::GuestError,
+                        format!("Unable to convert buffer to {}", type_name::<T>()),
+                    ));
+                }
             }
-        };
-        hyperlight_guest_tracing_macro::trace!("Finish converting buffer");
+        );
 
         // update the stack pointer to point to the element we just popped of since that is now free
         idb[..8].copy_from_slice(&last_element_offset_rel.to_le_bytes());
@@ -137,9 +138,9 @@ impl GuestHandle {
         }
 
         // write the actual data
-        hyperlight_guest_tracing_macro::trace!("Start copy of data");
-        odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(&data);
-        hyperlight_guest_tracing_macro::trace!("Finish copy of data");
+        hyperlight_guest_tracing_macro::trace!("copy data", {
+            odb[stack_ptr_rel as usize..stack_ptr_rel as usize + data.len()].copy_from_slice(&data);
+        });
 
         // write the offset to the newly written data, to the top of the stack
         let bytes: [u8; 8] = stack_ptr_rel.to_le_bytes();

src/hyperlight_guest_bin/src/exceptions/handler.rs

@@ -60,12 +60,15 @@ pub type HandlerT = fn(n: u64, info: *mut ExceptionInfo, ctx: *mut Context, pf_a
 
 /// Exception handler
 #[unsafe(no_mangle)]
-#[hyperlight_guest_tracing_macro::trace_function]
 pub extern "C" fn hl_exception_handler(
     stack_pointer: u64,
     exception_number: u64,
     page_fault_address: u64,
 ) {
+    // This function does not return, so wrapping it in a tracing macro
+    // issues a warning of unreachable code.
+    hyperlight_guest_tracing_macro::trace!("> hl_exception_handler");
+
     let ctx = stack_pointer as *mut Context;
     let exn_info = (stack_pointer + size_of::<Context>() as u64) as *mut ExceptionInfo;
 

src/hyperlight_guest_bin/src/guest_function/call.rs

@@ -62,8 +62,7 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
             core::mem::transmute::<usize, GuestFunc>(function_pointer)
         };
 
-        hyperlight_guest_tracing_macro::trace!("Calling guest function");
-        p_function(&function_call)
+        hyperlight_guest_tracing_macro::trace!("guest_function", p_function(&function_call))
     } else {
         // The given function is not registered. The guest should implement a function called guest_dispatch_function to handle this.
 
@@ -74,8 +73,9 @@ pub(crate) fn call_guest_function(function_call: FunctionCall) -> Result<Vec<u8>
             fn guest_dispatch_function(function_call: FunctionCall) -> Result<Vec<u8>>;
         }
 
-        hyperlight_guest_tracing_macro::trace!("Guest function {} not found");
-        unsafe { guest_dispatch_function(function_call) }
+        hyperlight_guest_tracing_macro::trace!("default guest function", unsafe {
+            guest_dispatch_function(function_call)
+        })
     }
 }
 

src/hyperlight_guest_bin/src/lib.rs

@@ -206,8 +206,9 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
 
             (*peb_ptr).guest_function_dispatch_ptr = dispatch_function as usize as u64;
 
-            trace!("hyperlight_main");
-            hyperlight_main();
+            trace!("hyperlight_main",
+                hyperlight_main();
+            );
         }
     });
 

src/hyperlight_guest_tracing/src/lib.rs

@@ -23,14 +23,18 @@ use core::mem::MaybeUninit;
 use hyperlight_common::outb::OutBAction;
 use spin::Mutex;
 
+/// Type alias for the function that sends trace records to the host.
+type SendToHostFn = fn(u64, &[TraceRecord]);
+
 /// Global trace buffer for storing trace records.
-static TRACE_BUFFER: Mutex<TraceBuffer> = Mutex::new(TraceBuffer::new());
+static TRACE_BUFFER: Mutex<TraceBuffer> = Mutex::new(TraceBuffer::new(send_to_host));
 
 /// Maximum number of entries in the trace buffer.
-const MAX_NO_OF_ENTRIES: usize = 64;
+/// From local testing, 32 entries seems to be a good balance between performance and memory usage.
+const MAX_NO_OF_ENTRIES: usize = 32;
 
 /// Maximum length of a trace message in bytes.
-pub const MAX_TRACE_MSG_LEN: usize = 64;
+const MAX_TRACE_MSG_LEN: usize = 64;
 
 #[derive(Debug, Copy, Clone)]
 /// Represents a trace record of a guest with a number of cycles and a message.
@@ -43,20 +47,45 @@ pub struct TraceRecord {
     pub msg: [u8; MAX_TRACE_MSG_LEN],
 }
 
+impl TryFrom<&str> for TraceRecord {
+    type Error = &'static str;
+
+    fn try_from(msg: &str) -> Result<Self, Self::Error> {
+        if msg.len() > MAX_TRACE_MSG_LEN {
+            return Err("Message too long");
+        }
+
+        let cycles = invariant_tsc::read_tsc();
+
+        Ok(TraceRecord {
+            cycles,
+            msg: {
+                let mut arr = [0u8; MAX_TRACE_MSG_LEN];
+                arr[..msg.len()].copy_from_slice(msg.as_bytes());
+                arr
+            },
+            msg_len: msg.len(),
+        })
+    }
+}
+
 /// A buffer for storing trace records.
 struct TraceBuffer {
     /// The entries in the trace buffer.
     entries: [TraceRecord; MAX_NO_OF_ENTRIES],
     /// The index where the next entry will be written.
     write_index: usize,
+    /// Function to send the trace records to the host.
+    send_to_host: SendToHostFn,
 }
 
 impl TraceBuffer {
     /// Creates a new `TraceBuffer` with uninitialized entries.
-    const fn new() -> Self {
+    const fn new(f: SendToHostFn) -> Self {
         Self {
             entries: unsafe { [MaybeUninit::zeroed().assume_init(); MAX_NO_OF_ENTRIES] },
             write_index: 0,
+            send_to_host: f,
         }
     }
 
@@ -72,26 +101,26 @@ impl TraceBuffer {
 
         if write_index == 0 {
             // If buffer is full send to host
-            self.send_to_host(MAX_NO_OF_ENTRIES);
+            (self.send_to_host)(MAX_NO_OF_ENTRIES as u64, &self.entries);
         }
     }
 
     /// Flush the trace buffer, sending any remaining records to the host.
     fn flush(&mut self) {
         if self.write_index > 0 {
-            self.send_to_host(self.write_index);
+            (self.send_to_host)(self.write_index as u64, &self.entries);
             self.write_index = 0; // Reset write index after flushing
         }
     }
+}
 
-    /// Send the trace records to the host.
-    fn send_to_host(&self, count: usize) {
-        unsafe {
-            core::arch::asm!("out dx, al",
+/// Send the trace records to the host.
+fn send_to_host(len: u64, records: &[TraceRecord]) {
+    unsafe {
+        core::arch::asm!("out dx, al",
                 in("dx") OutBAction::TraceRecord as u16,
-                in("rax") count as u64,
-                in("rcx") &self.entries as * const _ as u64);
-        }
+                in("rax") len,
+                in("rcx") records.as_ptr() as u64);
     }
 }
 
@@ -129,33 +158,120 @@ pub mod invariant_tsc {
     }
 }
 
-/// Create a trace record with the given message.
+/// Attempt to create a trace record from the message.
+/// If the message is too long it will skip the record.
+/// This is useful for ensuring that the trace buffer does not overflow.
+/// If the message is successfully converted, it will be pushed to the trace buffer.
 ///
-/// Note: The message must not exceed `MAX_TRACE_MSG_LEN` bytes.
-/// If the message is too long, it will be skipped.
+/// **Note**: The message must not exceed `MAX_TRACE_MSG_LEN` bytes.
 pub fn create_trace_record(msg: &str) {
-    if msg.len() > MAX_TRACE_MSG_LEN {
-        return; // Message too long, skip tracing
+    // Maybe panic if the message is too long?
+    if let Ok(entry) = TraceRecord::try_from(msg) {
+        let mut buffer = TRACE_BUFFER.lock();
+        buffer.push(entry);
     }
-
-    let cycles = invariant_tsc::read_tsc();
-
-    let entry = TraceRecord {
-        cycles,
-        msg: {
-            let mut arr = [0u8; MAX_TRACE_MSG_LEN];
-            arr[..msg.len()].copy_from_slice(msg.as_bytes());
-            arr
-        },
-        msg_len: msg.len(),
-    };
-
-    let mut buffer = TRACE_BUFFER.lock();
-    buffer.push(entry);
 }
 
 /// Flush the trace buffer to send any remaining trace records to the host.
 pub fn flush_trace_buffer() {
     let mut buffer = TRACE_BUFFER.lock();
     buffer.flush();
 }
+
+#[cfg(test)]
+mod tests {
+    use alloc::format;
+
+    use super::*;
+
+    /// This is a mock function for testing purposes.
+    /// In a real scenario, this would send the trace records to the host.
+    fn mock_send_to_host(_len: u64, _records: &[TraceRecord]) {}
+
+    fn create_test_entry(msg: &str) -> TraceRecord {
+        let cycles = invariant_tsc::read_tsc();
+
+        TraceRecord {
+            cycles,
+            msg: {
+                let mut arr = [0u8; MAX_TRACE_MSG_LEN];
+                arr[..msg.len()].copy_from_slice(msg.as_bytes());
+                arr
+            },
+            msg_len: msg.len(),
+        }
+    }
+
+    #[test]
+    fn test_push_trace_record() {
+        let mut buffer = TraceBuffer::new(mock_send_to_host);
+
+        let msg = "Test message";
+        let entry = create_test_entry(msg);
+
+        buffer.push(entry);
+        assert_eq!(buffer.write_index, 1);
+        assert_eq!(buffer.entries[0].msg_len, msg.len());
+        assert_eq!(&buffer.entries[0].msg[..msg.len()], msg.as_bytes());
+        assert!(buffer.entries[0].cycles > 0); // Ensure cycles is set
+    }
+
+    #[test]
+    fn test_flush_trace_buffer() {
+        let mut buffer = TraceBuffer::new(mock_send_to_host);
+
+        let msg = "Test message";
+        let entry = create_test_entry(msg);
+
+        buffer.push(entry);
+        assert_eq!(buffer.write_index, 1);
+        assert_eq!(buffer.entries[0].msg_len, msg.len());
+        assert_eq!(&buffer.entries[0].msg[..msg.len()], msg.as_bytes());
+        assert!(buffer.entries[0].cycles > 0);
+
+        // Flush the buffer
+        buffer.flush();
+
+        // After flushing, the entryes should still be intact, we don't clear them
+        assert_eq!(buffer.write_index, 0);
+        assert_eq!(buffer.entries[0].msg_len, msg.len());
+        assert_eq!(&buffer.entries[0].msg[..msg.len()], msg.as_bytes());
+        assert!(buffer.entries[0].cycles > 0);
+    }
+
+    #[test]
+    fn test_auto_flush_on_full() {
+        let mut buffer = TraceBuffer::new(mock_send_to_host);
+
+        // Fill the buffer to trigger auto-flush
+        for i in 0..MAX_NO_OF_ENTRIES {
+            let msg = format!("Message {}", i);
+            let entry = create_test_entry(&msg);
+            buffer.push(entry);
+        }
+
+        // After filling, the write index should be 0 (buffer is full)
+        assert_eq!(buffer.write_index, 0);
+
+        // The first entry should still be intact
+        assert_eq!(buffer.entries[0].msg_len, "Message 0".len());
+    }
+
+    /// Test TraceRecord creation with a valid message
+    #[test]
+    fn test_trace_record_creation_valid() {
+        let msg = "Valid message";
+        let entry = TraceRecord::try_from(msg).expect("Failed to create TraceRecord");
+        assert_eq!(entry.msg_len, msg.len());
+        assert_eq!(&entry.msg[..msg.len()], msg.as_bytes());
+        assert!(entry.cycles > 0); // Ensure cycles is set
+    }
+
+    /// Test TraceRecord creation with a message that exceeds the maximum length
+    #[test]
+    fn test_trace_record_creation_too_long() {
+        let long_msg = "A".repeat(MAX_TRACE_MSG_LEN + 1);
+        let result = TraceRecord::try_from(long_msg.as_str());
+        assert!(result.is_err(), "Expected error for message too long");
+    }
+}