Code refactor to set trace context in run_vcpu function. (#1184)

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

Author: vshailesh

src/hyperlight_host/src/hypervisor/hyperlight_vm.rs

@@ -27,8 +27,6 @@ use std::sync::{Arc, Mutex};
 
 use log::LevelFilter;
 use tracing::{Span, instrument};
-#[cfg(feature = "trace_guest")]
-use tracing_opentelemetry::OpenTelemetrySpanExt;
 
 #[cfg(gdb)]
 use super::gdb::arch::VcpuStopReasonError;
@@ -660,13 +658,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")]
@@ -134,7 +138,14 @@ 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> {
+        // setup_trace_guest must be called right before vcpu_run.run() call, because
+        // it sets the guest span, no other traces or spans must be setup in between these calls.
+        #[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

@@ -25,6 +25,8 @@ 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)]
@@ -286,8 +288,12 @@ pub(crate) trait VirtualMachine: Debug + Send {
     ) -> std::result::Result<(), UnmapMemoryError>;
 
     /// 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>;
+    /// Note: this function emits traces spans for guests
+    /// and the span setup is called right before the run virtual processor call of each hypervisor
+    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,10 @@ 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> {
         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;
@@ -130,6 +137,10 @@ impl VirtualMachine for MshvVm {
         #[cfg(gdb)]
         const EXCEPTION_INTERCEPT: hv_message_type = hv_message_type_HVMSG_X64_EXCEPTION_INTERCEPT;
 
+        // setup_trace_guest must be called right before vcpu_run.run() call, because
+        // it sets the guest span, no other traces or spans must be setup in between these calls.
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
         let exit_reason = self.vcpu_fd.run();
 
         let result = match exit_reason {

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

@@ -16,6 +16,9 @@ limitations under the License.
 
 use std::os::raw::c_void;
 
+use tracing::Span;
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use windows::Win32::Foundation::{FreeLibrary, HANDLE};
 use windows::Win32::System::Hypervisor::*;
 use windows::Win32::System::LibraryLoader::*;
@@ -36,6 +39,8 @@ use crate::hypervisor::virtual_machine::{
     VirtualMachine, VmExit,
 };
 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 {
@@ -213,9 +218,16 @@ 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> {
         let mut exit_context: WHV_RUN_VP_EXIT_CONTEXT = Default::default();
 
+        // setup_trace_guest must be called right before WHvRunVirtualProcessor() call, because
+        // it sets the guest span, no other traces or spans must be setup in between these calls.
+        #[cfg(feature = "trace_guest")]
+        tc.setup_guest_trace(Span::current().context());
         unsafe {
             WHvRunVirtualProcessor(
                 self.partition,
@@ -225,7 +237,6 @@ impl VirtualMachine for WhpVm {
             )
             .map_err(|e| RunVcpuError::Unknown(e.into()))?;
         }
-
         let result = match exit_context.ExitReason {
             WHvRunVpExitReasonX64IoPortAccess => unsafe {
                 let instruction_length = exit_context.VpContext._bitfield & 0xF;