[trace-host] refactor mem_profile logic

- Define a separate struct that holds the functionality related to
  memory profiling of the guest

Author: dblnz

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -320,6 +320,7 @@ pub(crate) struct HypervLinuxDriver {
     gdb_conn: Option<DebugCommChannel<DebugResponse, DebugMsg>>,
     #[cfg(crashdump)]
     rt_cfg: SandboxRuntimeConfig,
+    #[allow(dead_code)]
     #[cfg(feature = "trace_guest")]
     trace_info: TraceInfo,
 }
@@ -774,14 +775,6 @@ impl Hypervisor for HypervLinuxDriver {
         {
             Err(mshv_ioctls::MshvError::from(libc::EINTR))
         } else {
-            #[cfg(feature = "trace_guest")]
-            if self.trace_info.guest_start_epoch.is_none() {
-                // Store the guest start epoch and cycles to trace the guest execution time
-                crate::debug!("MSHV - Guest Start Epoch set");
-                self.trace_info.guest_start_tsc =
-                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
-                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
-            }
             // 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
@@ -1147,7 +1140,7 @@ impl Hypervisor for HypervLinuxDriver {
         Ok(X86_64Regs::from(self.vcpu_fd.get_regs()?))
     }
 
-    #[cfg(feature = "trace_guest")]
+    #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut TraceInfo {
         &mut self.trace_info
     }

src/hyperlight_host/src/hypervisor/hyperv_windows.rs

@@ -761,14 +761,6 @@ impl Hypervisor for HypervWindowsDriver {
                 Reserved: Default::default(),
             }
         } else {
-            #[cfg(feature = "trace_guest")]
-            if self.trace_info.guest_start_epoch.is_none() {
-                // Store the guest start epoch and cycles to trace the guest execution time
-                crate::debug!("HyperV - Guest Start Epoch set");
-                self.trace_info.guest_start_tsc =
-                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
-                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
-            }
             self.processor.run()?
         };
         self.interrupt_handle
@@ -1100,7 +1092,7 @@ impl Hypervisor for HypervWindowsDriver {
         Ok(X86_64Regs::from(regs))
     }
 
-    #[cfg(feature = "trace_guest")]
+    #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut TraceInfo {
         &mut self.trace_info
     }

src/hyperlight_host/src/hypervisor/kvm.rs

@@ -693,15 +693,6 @@ impl Hypervisor for KVMDriver {
         {
             Err(kvm_ioctls::Error::new(libc::EINTR))
         } else {
-            #[cfg(feature = "trace_guest")]
-            if self.trace_info.guest_start_epoch.is_none() {
-                // Store the guest start epoch and cycles to trace the guest execution time
-                crate::debug!("KVM - Guest Start Epoch set");
-                self.trace_info.guest_start_epoch = Some(std::time::Instant::now());
-                self.trace_info.guest_start_tsc =
-                    Some(hyperlight_guest_tracing::invariant_tsc::read_tsc());
-            }
-
             // 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
@@ -1041,7 +1032,7 @@ impl Hypervisor for KVMDriver {
         Ok(X86_64Regs::from(self.vcpu_fd.get_regs()?))
     }
 
-    #[cfg(feature = "trace_guest")]
+    #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut TraceInfo {
         &mut self.trace_info
     }

src/hyperlight_host/src/hypervisor/mod.rs

@@ -21,7 +21,7 @@ use crate::HyperlightError::StackOverflow;
 use crate::error::HyperlightError::ExecutionCanceledByHost;
 use crate::mem::memory_region::{MemoryRegion, MemoryRegionFlags};
 use crate::metrics::METRIC_GUEST_CANCELLATION;
-#[cfg(feature = "trace_guest")]
+#[cfg(feature = "mem_profile")]
 use crate::sandbox::trace::TraceInfo;
 use crate::{HyperlightError, Result, log_then_return};
 
@@ -233,11 +233,12 @@ pub(crate) trait Hypervisor: Debug + Send {
     fn check_stack_guard(&self) -> Result<bool>;
 
     /// Read a register for trace/unwind purposes
+    #[allow(dead_code)]
     #[cfg(feature = "trace_guest")]
     fn read_regs(&self) -> Result<arch::X86_64Regs>;
 
     /// Get a mutable reference of the trace info for the guest
-    #[cfg(feature = "trace_guest")]
+    #[cfg(feature = "mem_profile")]
     fn trace_info_mut(&mut self) -> &mut TraceInfo;
 }
 

src/hyperlight_host/src/sandbox/outb.rs

@@ -184,12 +184,14 @@ pub(crate) fn handle_outb(
         #[cfg(feature = "mem_profile")]
         OutBAction::TraceMemoryAlloc => {
             let regs = _hv.read_regs()?;
-            crate::sandbox::trace::handle_trace_memory_alloc(&regs, mem_mgr, _hv.trace_info_mut())
+            let trace_info = _hv.trace_info_mut();
+            trace_info.handle_trace_mem_alloc(&regs, mem_mgr.as_ref())
         }
         #[cfg(feature = "mem_profile")]
         OutBAction::TraceMemoryFree => {
             let regs = _hv.read_regs()?;
-            crate::sandbox::trace::handle_trace_memory_free(&regs, mem_mgr, _hv.trace_info_mut())
+            let trace_info = _hv.trace_info_mut();
+            trace_info.handle_trace_mem_free(&regs, mem_mgr.as_ref())
         }
     }
 }