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

Signed-off-by: Shailesh Vashishth <shavashishth@gmail.com>

Author: vshailesh

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -623,13 +623,13 @@ 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(
+                    #[cfg(feature = "trace_guest")]
+                    &mut 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

@@ -22,6 +22,8 @@ use kvm_bindings::{kvm_debugregs, kvm_fpu, kvm_regs, kvm_sregs, kvm_userspace_me
 use kvm_ioctls::Cap::UserMemory;
 use kvm_ioctls::{Kvm, VcpuExit, VcpuFd, VmFd};
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 #[cfg(gdb)]
 use crate::hypervisor::gdb::{DebugError, DebuggableVm};
@@ -33,6 +35,8 @@ use crate::hypervisor::virtual_machine::{
     VmExit,
 };
 use crate::mem::memory_region::MemoryRegion;
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::trace::TraceContext as SandboxTraceContext;
 
 /// Return `true` if the KVM API is available, version 12, and has UserMemory capability, or `false` otherwise
 #[instrument(skip_all, parent = Span::current(), level = "Trace")]
@@ -117,7 +121,13 @@ impl VirtualMachine for KvmVm {
             .map_err(|e| UnmapMemoryError::Hypervisor(e.into()))
     }
 
-    fn run_vcpu(&mut self) -> std::result::Result<VmExit, RunVcpuError> {
+    fn run_vcpu(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut SandboxTraceContext,
+    ) -> std::result::Result<VmExit, RunVcpuError> {
+        #[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

@@ -18,13 +18,17 @@ use std::fmt::Debug;
 use std::sync::OnceLock;
 
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 #[cfg(gdb)]
 use crate::hypervisor::gdb::DebugError;
 use crate::hypervisor::regs::{
     CommonDebugRegs, CommonFpu, CommonRegisters, CommonSpecialRegisters,
 };
 use crate::mem::memory_region::MemoryRegion;
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::trace::TraceContext as SandboxTraceContext;
 
 /// KVM (Kernel-based Virtual Machine) functionality (linux)
 #[cfg(kvm)]
@@ -287,7 +291,10 @@ 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) -> std::result::Result<VmExit, RunVcpuError>;
+    fn run_vcpu(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut SandboxTraceContext,
+    ) -> std::result::Result<VmExit, RunVcpuError>;
 
     /// Get regs
     #[allow(dead_code)]

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

@@ -30,6 +30,8 @@ use mshv_bindings::{
 };
 use mshv_ioctls::{Mshv, VcpuFd, VmFd};
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 #[cfg(gdb)]
 use crate::hypervisor::gdb::{DebugError, DebuggableVm};
@@ -41,6 +43,8 @@ use crate::hypervisor::virtual_machine::{
     VmExit,
 };
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::trace::TraceContext as SandboxTraceContext;
 
 /// Determine whether the HyperV for Linux hypervisor API is present
 /// and functional.
@@ -121,7 +125,12 @@ impl VirtualMachine for MshvVm {
             .map_err(|e| UnmapMemoryError::Hypervisor(e.into()))
     }
 
-    fn run_vcpu(&mut self) -> std::result::Result<VmExit, RunVcpuError> {
+    fn run_vcpu(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut SandboxTraceContext,
+    ) -> std::result::Result<VmExit, RunVcpuError> {
+        #[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

@@ -17,6 +17,8 @@ limitations under the License.
 use std::os::raw::c_void;
 
 use hyperlight_common::mem::PAGE_SIZE_USIZE;
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use windows::Win32::Foundation::{FreeLibrary, HANDLE};
 use windows::Win32::System::Hypervisor::*;
 use windows::Win32::System::LibraryLoader::*;
@@ -38,6 +40,8 @@ use crate::hypervisor::virtual_machine::{
 };
 use crate::hypervisor::wrappers::HandleWrapper;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
+#[cfg(feature = "trace_guest")]
+use crate::sandbox::trace::TraceContext as SandboxTraceContext;
 
 #[allow(dead_code)] // Will be used for runtime hypervisor detection
 pub(crate) fn is_hypervisor_present() -> bool {
@@ -239,7 +243,12 @@ impl VirtualMachine for WhpVm {
     }
 
     #[expect(non_upper_case_globals, reason = "Windows API constant are lower case")]
-    fn run_vcpu(&mut self) -> std::result::Result<VmExit, RunVcpuError> {
+    fn run_vcpu(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut SandboxTraceContext,
+    ) -> std::result::Result<VmExit, RunVcpuError> {
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
         let mut exit_context: WHV_RUN_VP_EXIT_CONTEXT = Default::default();
 
         unsafe {