[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,32 @@ 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.spans_ptr,
+                    in("r10") tbi.guest_start_tsc,
+                    options(nostack)
+                )
+            };
+        } 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 +58,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 +71,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);
@@ -81,6 +112,28 @@ 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.spans_ptr,
+                    in("r10") tbi.guest_start_tsc,
+                    options(preserves_flags, nomem, nostack)
+                )
+            };
+        } 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

@@ -16,7 +16,7 @@ limitations under the License.
 #![no_std]
 use heapless as hl;
 #[cfg(feature = "trace")]
-pub use trace::init_guest_tracing;
+pub use trace::{TraceBatchInfo, end_trace, guest_trace_info, init_guest_tracing};
 
 /// Module for checking invariant TSC support and reading the timestamp counter
 pub mod invariant_tsc {
@@ -139,18 +139,27 @@ mod trace {
     use alloc::sync::{Arc, Weak};
     use core::fmt::Debug;
     use core::sync::atomic::{AtomicU64, Ordering};
+
+    use hyperlight_common::outb::OutBAction;
+    use spin::Mutex;
     use tracing_core::field::{Field, Visit};
     use tracing_core::span::{Attributes, Id, Record};
     use tracing_core::subscriber::Subscriber;
     use tracing_core::{Event, Metadata};
-    use spin::Mutex;
 
     use super::*;
     use crate::invariant_tsc;
 
     /// Weak reference to the guest state so we can manually trigger flush to host
     static GUEST_STATE: spin::Once<Weak<Mutex<GuestState>>> = spin::Once::new();
 
+    pub struct TraceBatchInfo {
+        /// The timestamp counter at the start of the guest execution.
+        pub guest_start_tsc: u64,
+        /// Pointer to the spans in the guest memory.
+        pub spans_ptr: u64,
+    }
+
     /// Visitor implementation to collect fields into a vector of key-value pairs
     struct FieldsVisitor<'a, const FK: usize, const FV: usize, const F: usize> {
         out: &'a mut hl::Vec<(hl::String<FK>, hl::String<FV>), F>,
@@ -207,6 +216,8 @@ mod trace {
         const FV: usize,
         const F: usize,
     > {
+        /// Whether we need to cleanup the state on next access
+        cleanup_needed: bool,
         /// The timestamp counter at the start of the guest execution.
         guest_start_tsc: u64,
         /// Next span ID to allocate
@@ -229,6 +240,7 @@ mod trace {
     {
         fn new(guest_start_tsc: u64) -> Self {
             Self {
+                cleanup_needed: false,
                 guest_start_tsc,
                 next_id: AtomicU64::new(1),
                 spans: hl::Vec::new(),
@@ -243,13 +255,91 @@ mod trace {
             (n, Id::from_u64(n))
         }
 
+        /// Cleanup internal 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.
+        fn clean(&mut self) {
+            // used for computing the spans that need to be removed
+            let mut ids: hl::Vec<u64, SP> = self.spans.iter().map(|s| s.id).collect();
+
+            for id in self.stack.iter() {
+                let position = ids.iter().position(|s| *s == *id).unwrap();
+                // remove the span id that is contained in the stack
+                ids.remove(position);
+            }
+
+            // Remove the spans with the remaining ids
+            for id in ids.into_iter() {
+                let spans = &mut self.spans;
+                let position = spans.iter().position(|s| s.id == id).unwrap();
+                spans.remove(position);
+            }
+
+            // Remove the events from the remaining spans
+            for s in self.spans.iter_mut() {
+                s.events.clear();
+            }
+        }
+
+        #[inline(always)]
+        fn verify_and_clean(&mut self) {
+            if self.cleanup_needed {
+                self.clean();
+                self.cleanup_needed = false;
+            }
+        }
+
         /// Triggers a VM exit to flush the current spans to the host.
         /// This also clears the internal state to start fresh.
         fn send_to_host(&mut self) {
+            let guest_start_tsc = self.guest_start_tsc;
+            let spans_ptr = self.spans.as_ptr() as u64;
+
+            unsafe {
+                core::arch::asm!("out dx, al",
+                    // Port value for tracing
+                    in("dx") OutBAction::TraceBatch as u16,
+                    // Additional magic number to identify the action
+                    in("r8") OutBAction::TraceBatch as u64,
+                    in("r9") spans_ptr,
+                    in("r10") guest_start_tsc,
+                );
+            }
+
+            self.clean();
+        }
+
+        /// Closes the trace by ending all spans
+        /// NOTE: This expects an outb call to send the spans to the host.
+        fn end_trace(&mut self) {
+            for span in self.spans.iter_mut() {
+                if span.end_tsc.is_none() {
+                    span.end_tsc = Some(invariant_tsc::read_tsc());
+                }
+            }
+
+            // Empty the stack
+            while self.stack.pop().is_some() {
+                // Pop all remaining spans from the stack
+            }
+
+            // Mark for clearing when re-entering the VM because we might
+            // not enter on the same place as we exited (e.g. halt)
+            self.cleanup_needed = true;
+        }
+
+        /// Returns information about the information needed by the host to read the spans.
+        pub fn guest_trace_info(&mut self) -> TraceBatchInfo {
+            self.cleanup_needed = true;
+            TraceBatchInfo {
+                guest_start_tsc: self.guest_start_tsc,
+                spans_ptr: self.spans.as_ptr() as u64,
+            }
         }
 
         /// Create a new span and push it on the stack
         pub fn new_span(&mut self, attrs: &Attributes) -> Id {
+            self.verify_and_clean();
             let (idn, id) = self.alloc_id();
 
             let md = attrs.metadata();
@@ -293,6 +383,7 @@ mod trace {
 
         /// Record an event in the current span (top of the stack)
         pub fn event(&mut self, event: &Event<'_>) {
+            self.verify_and_clean();
             let stack = &mut self.stack;
             let parent_id = stack.last().copied().unwrap_or(0);
 
@@ -430,4 +521,29 @@ mod trace {
         // Set global dispatcher
         let _ = tracing_core::dispatcher::set_global_default(tracing_core::Dispatch::new(sub));
     }
+
+    /// 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();
+            }
+        }
+    }
+
+    /// 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

@@ -181,6 +181,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.read_regs()?;