Closes #1076: Code refactor to set trace context in run_vcpu function.

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -381,13 +381,10 @@ impl HyperlightVm {
             {
                 Ok(VmExit::Cancelled())
             } else {
-                #[cfg(feature = "trace_guest")]
-                tc.setup_guest_trace(Span::current().context());
-
                 // ==== KILL() TIMING POINT 3: Before calling run() ====
                 // If kill() is called and ran to completion BEFORE this line executes:
                 //    - Will still do a VM entry, but signals will be sent until VM exits
-                let result = self.vm.run_vcpu();
+                let result = self.vm.run_vcpu(&tc);
 
                 // End current host trace by closing the current span that captures traces
                 // happening when a guest exits and re-enters.

src/hyperlight_host/src/hypervisor/virtual_machine/kvm.rs

@@ -28,6 +28,7 @@ use crate::hypervisor::gdb::DebuggableVm;
 use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters};
 use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit};
 use crate::mem::memory_region::MemoryRegion;
+use crate::sandbox::trace::context::TraceContext as SandboxTraceContext;
 use crate::{Result, new_error};
 
 /// Return `true` if the KVM API is available, version 12, and has UserMemory capability, or `false` otherwise
@@ -104,7 +105,10 @@ impl VirtualMachine for KvmVm {
         Ok(())
     }
 
-    fn run_vcpu(&mut self) -> Result<VmExit> {
+    fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result<VmExit> {
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
+
         match self.vcpu_fd.run() {
             Ok(VcpuExit::Hlt) => Ok(VmExit::Halt()),
             Ok(VcpuExit::IoOut(port, data)) => Ok(VmExit::IoOut(port, data.to_vec())),

src/hyperlight_host/src/hypervisor/virtual_machine/mod.rs

@@ -22,6 +22,7 @@ use tracing::{Span, instrument};
 use crate::Result;
 use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters};
 use crate::mem::memory_region::MemoryRegion;
+use crate::sandbox::trace::context::TraceContext as SandboxTraceContext;
 
 /// KVM (Kernel-based Virtual Machine) functionality (linux)
 #[cfg(kvm)]
@@ -147,7 +148,7 @@ pub(crate) trait VirtualMachine: Debug + Send {
 
     /// Runs the vCPU until it exits.
     /// Note: this function should not emit any traces or spans as it is called after guest span is setup
-    fn run_vcpu(&mut self) -> Result<VmExit>;
+    fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result<VmExit>;
 
     /// Get regs
     #[allow(dead_code)]

src/hyperlight_host/src/hypervisor/virtual_machine/mshv.rs

@@ -36,6 +36,7 @@ use crate::hypervisor::gdb::DebuggableVm;
 use crate::hypervisor::regs::{CommonFpu, CommonRegisters, CommonSpecialRegisters};
 use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit};
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
+use crate::sandbox::trace::context::TraceContext as SandboxTraceContext;
 use crate::{Result, new_error};
 
 /// Determine whether the HyperV for Linux hypervisor API is present
@@ -103,7 +104,9 @@ impl VirtualMachine for MshvVm {
         Ok(())
     }
 
-    fn run_vcpu(&mut self) -> Result<VmExit> {
+    fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result<VmExit> {
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
         const HALT_MESSAGE: hv_message_type = hv_message_type_HVMSG_X64_HALT;
         const IO_PORT_INTERCEPT_MESSAGE: hv_message_type =
             hv_message_type_HVMSG_X64_IO_PORT_INTERCEPT;

src/hyperlight_host/src/hypervisor/virtual_machine/whp.rs

@@ -34,6 +34,7 @@ use crate::hypervisor::surrogate_process_manager::get_surrogate_process_manager;
 use crate::hypervisor::virtual_machine::{VirtualMachine, VmExit};
 use crate::hypervisor::wrappers::HandleWrapper;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
+use crate::sandbox::trace::context::TraceContext as SandboxTraceContext;
 use crate::{Result, log_then_return, new_error};
 
 #[allow(dead_code)] // Will be used for runtime hypervisor detection
@@ -205,7 +206,9 @@ impl VirtualMachine for WhpVm {
     }
 
     #[expect(non_upper_case_globals, reason = "Windows API constant are lower case")]
-    fn run_vcpu(&mut self) -> Result<VmExit> {
+    fn run_vcpu(&mut self, tc: &SandboxTraceContext) -> Result<VmExit> {
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
         let mut exit_context: WHV_RUN_VP_EXIT_CONTEXT = Default::default();
 
         unsafe {