[trace-host] refactor mem_profile logic

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

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

Author: dblnz

src/hyperlight_host/src/hypervisor/hyperv_linux.rs

@@ -296,6 +296,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,
 }
@@ -684,14 +685,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
@@ -1093,7 +1086,7 @@ impl Hypervisor for HypervLinuxDriver {
         }
     }
 
-    #[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

@@ -569,14 +569,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
@@ -940,7 +932,7 @@ impl Hypervisor for HypervWindowsDriver {
         }
     }
 
-    #[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

@@ -646,15 +646,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
@@ -1030,7 +1021,7 @@ impl Hypervisor for KVMDriver {
         }
     }
 
-    #[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

@@ -24,7 +24,7 @@ use crate::hypervisor::regs::{
 };
 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};
 
@@ -314,7 +314,7 @@ pub(crate) trait Hypervisor: Debug + Send {
     fn check_stack_guard(&self) -> Result<bool>;
 
     /// 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

@@ -185,12 +185,14 @@ pub(crate) fn handle_outb(
         #[cfg(feature = "mem_profile")]
         OutBAction::TraceMemoryAlloc => {
             let regs = _hv.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)
         }
         #[cfg(feature = "mem_profile")]
         OutBAction::TraceMemoryFree => {
             let regs = _hv.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)
         }
     }
 }