fix(trace): read each guest trace batch once

docs/hyperlight-metrics-logs-and-traces.md

@@ -152,11 +152,11 @@ applies to later guest calls and is reapplied after snapshot restore.
 When the storage space is filled, the guest triggers a VM Exit that sends the guest pointers to the host. The host can access the guest memory, get the data and parse it to create the `spans` and `events` using the `opentelemetry` crate which allows specifying the starting and ending timestamps
 which are captured in the guest using the `TSC`.
 
-To improve performance, for each VMExit, the guest adds metadata for the host to be able to report the tracing data and free space.
+To improve performance, each OUT exit also carries any pending trace data. The guest sets `r8` to the `TraceBatch` marker, with the buffer pointer in `r9` and its length in `r10`. It clears `r8` after the exit, so each batch is reported once.
 
 ##### Host
 
-When a guest exits, the host checks for metadata from the guest reporting tracing data.
+On each OUT exit, the host checks `r8` for the `TraceBatch` marker. Other exits carry no trace data.
 If tracing data is found, the host starts parsing it and reconstructing a tree which represents the spans hierarchy.
 
 Additionally, the host also adds new children `span`s to the guest's reported active span, emphasizing the spans created on the host as a result of a temporary VM Exit. This helps visualize a call into the guest with context propagated across the VM boundary.

src/hyperlight_guest/src/arch/amd64/exit.rs

@@ -17,12 +17,14 @@ pub(crate) unsafe fn out32(port: u16, val: u32) {
     #[cfg(feature = "trace_guest")]
     {
         if let Some((ptr, len)) = hyperlight_guest_tracing::serialized_data() {
-            // If tracing is enabled and there is data to send, send it along with the OUT action
+            // If tracing is enabled and there is data to send, send it along with the OUT action.
+            // Clearing r8 after the exit keeps later exits from reporting this batch again.
             unsafe {
                 asm!("out dx, eax",
+                    "mov r8d, 0",
                     in("dx") port,
                     in("eax") val,
-                    in("r8") OutBAction::TraceBatch as u64,
+                    inout("r8") OutBAction::TraceBatch as u64 => _,
                     in("r9") ptr,
                     in("r10") len,
                     options(preserves_flags, nostack)

src/hyperlight_guest_tracing/src/state.rs

@@ -37,11 +37,13 @@ const ACTIVE_SPANS_CAPACITY: usize = 64;
 fn send_to_host(data: &[u8]) {
     unsafe {
         core::arch::asm!("out dx, al",
+            // Clear the marker so later exits do not report this batch again
+            "mov r8d, 0",
             // Port value for tracing
             in("dx") OutBAction::TraceBatch as u16,
             in("al") 0u8,
             // Additional magic number to identify the action
-            in("r8") OutBAction::TraceBatch as u64,
+            inout("r8") OutBAction::TraceBatch as u64 => _,
             in("r9") data.as_ptr() as u64,
             in("r10") data.len() as u64,
         );

src/hyperlight_host/src/hypervisor/hyperlight_vm/mod.rs

@@ -620,20 +620,23 @@ impl HyperlightVm {
                 #[cfg(feature = "trace_guest")]
                 {
                     tc.end_host_trace();
-                    // Handle the guest trace data if any
-                    let regs = self.vm.regs().map_err(RunVmError::GetRegs)?;
-
-                    // Only parse the trace if it has reported
-                    if tc.has_trace_data(&regs) {
-                        let root_pt = self.get_root_pt().map_err(RunVmError::PageTableAccess)?;
-
-                        // If something goes wrong with parsing the trace data, we log the error and
-                        // continue execution instead of returning an error since this is not critical
-                        // to correct execution of the guest
-                        tc.handle_trace(&regs, mem_mgr, root_pt)
-                            .unwrap_or_else(|e| {
-                                tracing::error!("Cannot handle trace data: {}", e);
-                            });
+                    // The guest reports trace batches only on OUT exits. Other exits
+                    // can hold stale registers and must not be checked.
+                    if let Ok(VmExit::IoOut(..)) = &result {
+                        let regs = self.vm.regs().map_err(RunVmError::GetRegs)?;
+
+                        if tc.has_trace_data(&regs) {
+                            let root_pt =
+                                self.get_root_pt().map_err(RunVmError::PageTableAccess)?;
+
+                            // If something goes wrong with parsing the trace data, we log the error and
+                            // continue execution instead of returning an error since this is not critical
+                            // to correct execution of the guest
+                            tc.handle_trace(&regs, mem_mgr, root_pt)
+                                .unwrap_or_else(|e| {
+                                    tracing::error!("Cannot handle trace data: {}", e);
+                                });
+                        }
                     }
                 }
                 result