[trace-guest] update outb instructions to include trace info

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

Author: dblnz

src/hyperlight_common/src/outb.rs

@@ -90,13 +90,16 @@ impl TryFrom<u8> for Exception {
 /// - CallFunction: makes a call to a host function,
 /// - Abort: aborts the execution of the guest,
 /// - DebugPrint: prints a message to the host
+/// - TraceBatch: reports a batch of spans and vents from the guest
 /// - TraceMemoryAlloc: records memory allocation events
 /// - TraceMemoryFree: records memory deallocation events
 pub enum OutBAction {
     Log = 99,
     CallFunction = 101,
     Abort = 102,
     DebugPrint = 103,
+    #[cfg(feature = "trace_guest")]
+    TraceBatch = 104,
     #[cfg(feature = "mem_profile")]
     TraceMemoryAlloc = 105,
     #[cfg(feature = "mem_profile")]
@@ -111,6 +114,8 @@ impl TryFrom<u16> for OutBAction {
             101 => Ok(OutBAction::CallFunction),
             102 => Ok(OutBAction::Abort),
             103 => Ok(OutBAction::DebugPrint),
+            #[cfg(feature = "trace_guest")]
+            104 => Ok(OutBAction::TraceBatch),
             #[cfg(feature = "mem_profile")]
             105 => Ok(OutBAction::TraceMemoryAlloc),
             #[cfg(feature = "mem_profile")]

src/hyperlight_guest/src/exit.rs

@@ -22,7 +22,34 @@ use hyperlight_common::outb::OutBAction;
 /// Halt the execution of the guest and returns control to the host.
 #[inline(never)]
 pub fn halt() {
-    unsafe { asm!("hlt", options(nostack)) }
+    #[cfg(feature = "trace_guest")]
+    {
+        // End any ongoing trace before halting
+        hyperlight_guest_tracing::end_trace();
+        // If tracing is enabled, we need to pass the trace batch info
+        // along with the halt instruction so the host can retrieve it
+        if let Some(tbi) = hyperlight_guest_tracing::guest_trace_info() {
+            unsafe {
+                asm!("hlt",
+                    in("r8") OutBAction::TraceBatch as u64,
+                    in("r9") tbi.guest_start_tsc,
+                    in("r10") tbi.spans_ptr,
+                    in("r11") tbi.events_ptr,
+                    options(nostack)
+                )
+            };
+            hyperlight_guest_tracing::clean_trace_state();
+        } else {
+            // If tracing is not enabled, we can directly halt
+            unsafe { asm!("hlt", options(nostack)) };
+        }
+    }
+
+    #[cfg(not(feature = "trace_guest"))]
+    {
+        // If tracing is not enabled, we can directly halt
+        unsafe { asm!("hlt", options(nostack)) };
+    }
 }
 
 /// Exits the VM with an Abort OUT action and code 0.
@@ -33,6 +60,9 @@ pub extern "C" fn abort() -> ! {
 
 /// Exits the VM with an Abort OUT action and a specific code.
 pub fn abort_with_code(code: &[u8]) -> ! {
+    // End any ongoing trace before aborting
+    #[cfg(feature = "trace_guest")]
+    hyperlight_guest_tracing::end_trace();
     outb(OutBAction::Abort as u16, code);
     outb(OutBAction::Abort as u16, &[0xFF]); // send abort terminator (if not included in code)
     unreachable!()
@@ -43,6 +73,9 @@ pub fn abort_with_code(code: &[u8]) -> ! {
 /// # Safety
 /// This function is unsafe because it dereferences a raw pointer.
 pub unsafe fn abort_with_code_and_message(code: &[u8], message_ptr: *const c_char) -> ! {
+    // End any ongoing trace before aborting
+    #[cfg(feature = "trace_guest")]
+    hyperlight_guest_tracing::end_trace();
     unsafe {
         // Step 1: Send abort code (typically 1 byte, but `code` allows flexibility)
         outb(OutBAction::Abort as u16, code);
@@ -89,6 +122,31 @@ pub(crate) fn outb(port: u16, data: &[u8]) {
 
 /// OUT function for sending a 32-bit value to the host.
 pub(crate) unsafe fn out32(port: u16, val: u32) {
+    #[cfg(feature = "trace_guest")]
+    {
+        if let Some(tbi) = hyperlight_guest_tracing::guest_trace_info() {
+            // If tracing is enabled, send the trace batch info along with the OUT action
+            unsafe {
+                asm!("out dx, eax",
+                    in("dx") port,
+                    in("eax") val,
+                    in("r8") OutBAction::TraceBatch as u64,
+                    in("r9") tbi.guest_start_tsc,
+                    in("r10") tbi.spans_ptr,
+                    in("r11") tbi.events_ptr,
+                    options(preserves_flags, nomem, nostack)
+                )
+            };
+
+            hyperlight_guest_tracing::clean_trace_state();
+        } else {
+            // If tracing is not enabled, just send the value
+            unsafe {
+                asm!("out dx, eax", in("dx") port, in("eax") val, options(preserves_flags, nomem, nostack))
+            };
+        }
+    }
+    #[cfg(not(feature = "trace_guest"))]
     unsafe {
         asm!("out dx, eax", in("dx") port, in("eax") val, options(preserves_flags, nomem, nostack));
     }

src/hyperlight_guest_tracing/src/lib.rs

@@ -37,7 +37,9 @@ mod visitor;
 #[cfg(feature = "trace")]
 pub use state::TraceBatchInfo;
 #[cfg(feature = "trace")]
-pub use trace::{init_guest_tracing, set_start_tsc};
+pub use trace::{
+    clean_trace_state, end_trace, guest_trace_info, init_guest_tracing, set_start_tsc,
+};
 
 /// Maximum number of spans that the guest can store
 const MAX_NO_OF_SPANS: usize = 10;
@@ -174,4 +176,40 @@ mod trace {
             }
         }
     }
+
+    /// Ends the current trace by ending all active spans in the
+    /// internal state and storing the end timestamps.
+    ///
+    /// This expects an outb call to send the spans to the host.
+    /// After calling this function, the internal state is marked
+    /// for cleaning on the next access.
+    pub fn end_trace() {
+        if let Some(w) = GUEST_STATE.get() {
+            if let Some(state) = w.upgrade() {
+                state.lock().end_trace();
+            }
+        }
+    }
+
+    /// Cleans the internal trace state by removing closed spans and events.
+    /// This ensures that after a VM exit, we keep the spans that
+    /// are still active (in the stack) and remove all other spans and events.
+    pub fn clean_trace_state() {
+        if let Some(w) = GUEST_STATE.get() {
+            if let Some(state) = w.upgrade() {
+                state.lock().clean();
+            }
+        }
+    }
+
+    /// Returns information about the current trace state needed by the host to read the spans.
+    pub fn guest_trace_info() -> Option<TraceBatchInfo> {
+        let mut res = None;
+        if let Some(w) = GUEST_STATE.get() {
+            if let Some(state) = w.upgrade() {
+                res = Some(state.lock().guest_trace_info());
+            }
+        }
+        res
+    }
 }

src/hyperlight_host/src/sandbox/outb.rs

@@ -182,6 +182,8 @@ pub(crate) fn handle_outb(
             eprint!("{}", ch);
             Ok(())
         }
+        #[cfg(feature = "trace_guest")]
+        OutBAction::TraceBatch => Ok(()),
         #[cfg(feature = "mem_profile")]
         OutBAction::TraceMemoryAlloc => {
             let regs = _hv.regs()?;