[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

@@ -170,10 +170,10 @@ mod trace {
 
     /// Sets the guset starting timestamp reported to the host on a VMExit
     pub fn set_start_tsc(guest_start_tsc: u64) {
-        if let Some(w) = GUEST_STATE.get() {
-            if let Some(state) = w.upgrade() {
-                state.lock().set_start_tsc(guest_start_tsc);
-            }
+        if let Some(w) = GUEST_STATE.get()
+            && let Some(state) = w.upgrade()
+        {
+            state.lock().set_start_tsc(guest_start_tsc);
         }
     }
 
@@ -184,31 +184,31 @@ mod trace {
     /// After calling this function, the internal state is marked
     /// for cleaning on the next access.
     pub fn end_trace() {
-        if let Some(w) = GUEST_STATE.get() {
-            if let Some(state) = w.upgrade() {
-                state.lock().end_trace();
-            }
+        if let Some(w) = GUEST_STATE.get()
+            && let Some(state) = w.upgrade()
+        {
+            state.lock().end_trace();
         }
     }
 
     /// Cleans the internal trace state by removing closed spans and events.
     /// This ensures that after a VM exit, we keep the spans that
     /// are still active (in the stack) and remove all other spans and events.
     pub fn clean_trace_state() {
-        if let Some(w) = GUEST_STATE.get() {
-            if let Some(state) = w.upgrade() {
-                state.lock().clean();
-            }
+        if let Some(w) = GUEST_STATE.get()
+            && let Some(state) = w.upgrade()
+        {
+            state.lock().clean();
         }
     }
 
     /// Returns information about the current trace state needed by the host to read the spans.
     pub fn guest_trace_info() -> Option<TraceBatchInfo> {
         let mut res = None;
-        if let Some(w) = GUEST_STATE.get() {
-            if let Some(state) = w.upgrade() {
-                res = Some(state.lock().guest_trace_info());
-            }
+        if let Some(w) = GUEST_STATE.get()
+            && let Some(state) = w.upgrade()
+        {
+            res = Some(state.lock().guest_trace_info());
         }
         res
     }

src/hyperlight_guest_tracing/src/lib.rs

@@ -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"
@@ -98,7 +100,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.31"
 opentelemetry_sdk = { version = "0.30.0", features = ["rt-tokio"] }
 tokio = { version = "1.47.1", features = ["full"] }
@@ -132,7 +134,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};
 
@@ -319,7 +321,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,
 }
@@ -732,7 +733,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;
@@ -774,6 +778,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
@@ -1139,6 +1146,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,
@@ -292,7 +294,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
@@ -735,7 +736,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))]
@@ -760,6 +764,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
@@ -1091,6 +1098,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};
 
@@ -305,7 +307,6 @@ pub(crate) struct KVMDriver {
     #[cfg(crashdump)]
     rt_cfg: SandboxRuntimeConfig,
     #[cfg(feature = "mem_profile")]
-    #[allow(dead_code)]
     trace_info: MemTraceInfo,
 }
 
@@ -662,7 +663,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);
@@ -695,6 +699,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
@@ -1034,6 +1041,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