[trace-host] handle a TraceBatch from the guest

- Parse the spans and events coming from the guest and
create corresponding spans and events from the host that
mimics a single call from host
- Create a `TraceContext` that handles a call into a guest

Signed-off-by: Doru Blânzeanu <[email protected]>

Author: dblnz

Cargo.lock

@@ -35,9 +35,11 @@ page_size = "0.6.0"
 termcolor = "1.2.0"
 bitflags = "2.9.4"
 log = "0.4.28"
+opentelemetry = { version = "0.30.0", optional = true }
 tracing = { version = "0.1.41", features = ["log"] }
 tracing-log = "0.2.0"
 tracing-core = "0.1.34"
+tracing-opentelemetry = { version = "0.31.0", optional = true }
 hyperlight-common = { workspace = true, default-features = true, features = [ "std" ] }
 hyperlight-guest-tracing = { workspace = true, default-features = true, optional = true }
 vmm-sys-util = "0.15.0"
@@ -97,7 +99,7 @@ tracing = "0.1.41"
 tracing-subscriber = {version = "0.3.20", features = ["std", "env-filter"]}
 tracing-opentelemetry = "0.31.0"
 opentelemetry = "0.30.0"
-opentelemetry-otlp = { version = "0.30.0", default-features = false, features = ["http-proto", "reqwest-blocking-client"] }
+opentelemetry-otlp = { version = "0.30.0", default-features = false, features = ["http-proto", "reqwest-blocking-client", "grpc-tonic"] }
 opentelemetry-semantic-conventions = "0.30"
 opentelemetry_sdk = { version = "0.30.0", features = ["rt-tokio"] }
 tokio = { version = "1.47.1", features = ["full"] }
@@ -131,7 +133,7 @@ executable_heap = []
 print_debug = []
 # Dumps the VM state to a file on unexpected errors or crashes. The path of the file will be printed on stdout and logged.
 crashdump = ["dep:chrono"]
-trace_guest = ["hyperlight-common/trace_guest", "hyperlight-guest-tracing/trace"]
+trace_guest = ["dep:opentelemetry", "dep:tracing-opentelemetry", "dep:hyperlight-guest-tracing", "hyperlight-common/trace_guest"]
 mem_profile = [ "trace_guest", "dep:framehop", "dep:fallible-iterator", "hyperlight-common/mem_profile" ]
 kvm = ["dep:kvm-bindings", "dep:kvm-ioctls"]
 # This feature is deprecated in favor of mshv3

src/hyperlight_host/Cargo.toml

@@ -50,6 +50,8 @@ use mshv_bindings::{
 };
 use mshv_ioctls::{Mshv, VcpuFd, VmFd};
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 #[cfg(crashdump)]
 use {super::crashdump, std::path::Path};
 
@@ -320,7 +322,6 @@ pub(crate) struct HypervLinuxDriver {
     gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
     #[cfg(crashdump)]
     rt_cfg: SandboxRuntimeConfig,
-    #[allow(dead_code)]
     #[cfg(feature = "mem_profile")]
     trace_info: MemTraceInfo,
 }
@@ -733,7 +734,10 @@ impl Hypervisor for HypervLinuxDriver {
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    fn run(&mut self) -> Result<super::HyperlightExit> {
+    fn run(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut crate::sandbox::trace::TraceContext,
+    ) -> Result<super::HyperlightExit> {
         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;
@@ -775,6 +779,9 @@ impl Hypervisor for HypervLinuxDriver {
         {
             Err(mshv_ioctls::MshvError::from(libc::EINTR))
         } else {
+            #[cfg(feature = "trace_guest")]
+            tc.setup_guest_trace(Span::current().context());
+
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**
             // Then the vcpu will run, but we will keep sending signals to this thread
             // to interrupt it until `running` is set to false. The `vcpu_fd::run()` call will
@@ -1140,6 +1147,17 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(X86_64Regs::from(self.vcpu_fd.get_regs()?))
     }
 
+    #[cfg(feature = "trace_guest")]
+    fn handle_trace(&mut self, tc: &mut crate::sandbox::trace::TraceContext) -> Result<()> {
+        let regs = self.read_regs()?;
+        tc.handle_trace(
+            &regs,
+            self.mem_mgr.as_ref().ok_or_else(|| {
+                new_error!("Memory manager is not initialized before handling trace")
+            })?,
+        )
+    }
+
     #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut MemTraceInfo {
         &mut self.trace_info

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -22,6 +22,8 @@ use std::sync::{Arc, Mutex};
 
 use log::LevelFilter;
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 use windows::Win32::System::Hypervisor::{
     WHV_MEMORY_ACCESS_TYPE, WHV_PARTITION_HANDLE, WHV_REGISTER_VALUE, WHV_RUN_VP_EXIT_CONTEXT,
     WHV_RUN_VP_EXIT_REASON, WHV_X64_SEGMENT_REGISTER, WHV_X64_SEGMENT_REGISTER_0,
@@ -293,7 +295,6 @@ pub(crate) struct HypervWindowsDriver {
     #[cfg(crashdump)]
     rt_cfg: SandboxRuntimeConfig,
     #[cfg(feature = "mem_profile")]
-    #[allow(dead_code)]
     trace_info: MemTraceInfo,
 }
 /* This does not automatically impl Send because the host
@@ -736,7 +737,10 @@ impl Hypervisor for HypervWindowsDriver {
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    fn run(&mut self) -> Result<super::HyperlightExit> {
+    fn run(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut crate::sandbox::trace::TraceContext,
+    ) -> Result<super::HyperlightExit> {
         self.interrupt_handle.running.store(true, Ordering::Relaxed);
 
         #[cfg(not(gdb))]
@@ -761,6 +765,9 @@ impl Hypervisor for HypervWindowsDriver {
                 Reserved: Default::default(),
             }
         } else {
+            #[cfg(feature = "trace_guest")]
+            tc.setup_guest_trace(Span::current().context());
+
             self.processor.run()?
         };
         self.interrupt_handle
@@ -1092,6 +1099,17 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(X86_64Regs::from(regs))
     }
 
+    #[cfg(feature = "trace_guest")]
+    fn handle_trace(&mut self, tc: &mut crate::sandbox::trace::TraceContext) -> Result<()> {
+        let regs = self.read_regs()?;
+        tc.handle_trace(
+            &regs,
+            self.mem_mgr.as_ref().ok_or_else(|| {
+                new_error!("Memory manager is not initialized before handling trace")
+            })?,
+        )
+    }
+
     #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut MemTraceInfo {
         &mut self.trace_info

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -24,6 +24,8 @@ use kvm_ioctls::Cap::UserMemory;
 use kvm_ioctls::{Kvm, VcpuExit, VcpuFd, VmFd};
 use log::LevelFilter;
 use tracing::{Span, instrument};
+#[cfg(feature = "trace_guest")]
+use tracing_opentelemetry::OpenTelemetrySpanExt;
 #[cfg(crashdump)]
 use {super::crashdump, std::path::Path};
 
@@ -306,7 +308,6 @@ pub(crate) struct KVMDriver {
     #[cfg(crashdump)]
     rt_cfg: SandboxRuntimeConfig,
     #[cfg(feature = "mem_profile")]
-    #[allow(dead_code)]
     trace_info: MemTraceInfo,
 }
 
@@ -660,7 +661,10 @@ impl Hypervisor for KVMDriver {
     }
 
     #[instrument(err(Debug), skip_all, parent = Span::current(), level = "Trace")]
-    fn run(&mut self) -> Result<HyperlightExit> {
+    fn run(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut crate::sandbox::trace::TraceContext,
+    ) -> Result<HyperlightExit> {
         self.interrupt_handle
             .tid
             .store(unsafe { libc::pthread_self() as u64 }, Ordering::Relaxed);
@@ -693,6 +697,9 @@ impl Hypervisor for KVMDriver {
         {
             Err(kvm_ioctls::Error::new(libc::EINTR))
         } else {
+            #[cfg(feature = "trace_guest")]
+            tc.setup_guest_trace(Span::current().context());
+
             // Note: if a `InterruptHandle::kill()` called while this thread is **here**
             // Then the vcpu will run, but we will keep sending signals to this thread
             // to interrupt it until `running` is set to false. The `vcpu_fd::run()` call will
@@ -1032,6 +1039,17 @@ impl Hypervisor for KVMDriver {
         Ok(X86_64Regs::from(self.vcpu_fd.get_regs()?))
     }
 
+    #[cfg(feature = "trace_guest")]
+    fn handle_trace(&mut self, tc: &mut crate::sandbox::trace::TraceContext) -> Result<()> {
+        let regs = self.read_regs()?;
+        tc.handle_trace(
+            &regs,
+            self.mem_mgr.as_ref().ok_or_else(|| {
+                new_error!("Memory manager is not initialized before handling trace")
+            })?,
+        )
+    }
+
     #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut MemTraceInfo {
         &mut self.trace_info

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -171,7 +171,10 @@ pub(crate) trait Hypervisor: Debug + Send {
     ) -> Result<()>;
 
     /// Run the vCPU
-    fn run(&mut self) -> Result<HyperlightExit>;
+    fn run(
+        &mut self,
+        #[cfg(feature = "trace_guest")] tc: &mut crate::sandbox::trace::TraceContext,
+    ) -> Result<HyperlightExit>;
 
     /// Get InterruptHandle to underlying VM
     fn interrupt_handle(&self) -> Arc<dyn InterruptHandle>;
@@ -237,6 +240,9 @@ pub(crate) trait Hypervisor: Debug + Send {
     #[cfg(feature = "trace_guest")]
     fn read_regs(&self) -> Result<arch::X86_64Regs>;
 
+    #[cfg(feature = "trace_guest")]
+    fn handle_trace(&mut self, tc: &mut crate::sandbox::trace::TraceContext) -> Result<()>;
+
     /// Get a mutable reference of the trace info for the guest
     #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut MemTraceInfo;
@@ -276,8 +282,31 @@ impl VirtualCPU {
         hv: &mut dyn Hypervisor,
         #[cfg(gdb)] dbg_mem_access_fn: Arc<Mutex<MemMgrWrapper<HostSharedMemory>>>,
     ) -> Result<()> {
+        // Keeps the trace context and open spans
+        #[cfg(feature = "trace_guest")]
+        let mut tc = crate::sandbox::trace::TraceContext::new();
+
         loop {
-            match hv.run() {
+            #[cfg(feature = "trace_guest")]
+            let result = {
+                let result = hv.run(&mut tc);
+                // End current host trace by closing the current span that captures traces
+                // happening when a guest exits and re-enters.
+                tc.end_host_trace();
+
+                // Handle the guest trace data if any
+                if let Err(e) = hv.handle_trace(&mut tc) {
+                    // If no trace data is available, we just log a message and continue
+                    // Is this the right thing to do?
+                    log::debug!("Error handling guest trace: {:?}", e);
+                }
+
+                result
+            };
+            #[cfg(not(feature = "trace_guest"))]
+            let result = hv.run();
+
+            match result {
                 #[cfg(gdb)]
                 Ok(HyperlightExit::Debug(stop_reason)) => {
                     if let Err(e) = hv.handle_debug(dbg_mem_access_fn.clone(), stop_reason) {