Add trace support for recording memory allocations and frees

This allows for producing profiles of memory usage.

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

Author: syntactically

src/hyperlight_common/Cargo.toml

@@ -26,6 +26,7 @@ spin = "0.10.0"
 default = ["tracing"]
 fuzzing = ["dep:arbitrary"]
 unwind_guest = []
+mem_profile = []
 std = []
 
 [dev-dependencies]

src/hyperlight_common/src/outb.rs

@@ -91,13 +91,19 @@ impl TryFrom<u8> for Exception {
 /// - Abort: aborts the execution of the guest,
 /// - DebugPrint: prints a message to the host
 /// - TraceRecordStack: records the stack trace of the guest
+/// - TraceMemoryAlloc: records memory allocation events
+/// - TraceMemoryFree: records memory deallocation events
 pub enum OutBAction {
     Log = 99,
     CallFunction = 101,
     Abort = 102,
     DebugPrint = 103,
     #[cfg(feature = "unwind_guest")]
     TraceRecordStack = 104,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryAlloc,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryFree,
 }
 
 impl TryFrom<u16> for OutBAction {
@@ -110,6 +116,10 @@ impl TryFrom<u16> for OutBAction {
             103 => Ok(OutBAction::DebugPrint),
             #[cfg(feature = "unwind_guest")]
             104 => Ok(OutBAction::TraceRecordStack),
+            #[cfg(feature = "mem_profile")]
+            105 => Ok(OutBAction::TraceMemoryAlloc),
+            #[cfg(feature = "mem_profile")]
+            106 => Ok(OutBAction::TraceMemoryFree),
             _ => Err(anyhow::anyhow!("Invalid OutBAction value: {}", val)),
         }
     }

src/hyperlight_guest_bin/Cargo.toml

@@ -17,6 +17,7 @@ and third-party code used by our C-API needed to build a native hyperlight-guest
 default = ["libc", "printf"]
 libc = [] # compile musl libc
 printf = [] # compile printf
+mem_profile = ["hyperlight-common/unwind_guest","hyperlight-common/mem_profile"]
 
 [dependencies]
 hyperlight-guest = { workspace = true, default-features = false }

src/hyperlight_guest_bin/src/lib.rs

@@ -28,6 +28,8 @@ use guest_function::register::GuestFunctionRegister;
 use guest_logger::init_logger;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::mem::HyperlightPEB;
+#[cfg(feature = "mem_profile")]
+use hyperlight_common::outb::OutBAction;
 use hyperlight_guest::exit::{abort_with_code_and_message, halt};
 use hyperlight_guest::guest_handle::handle::GuestHandle;
 use log::LevelFilter;
@@ -53,9 +55,69 @@ pub mod guest_logger;
 pub mod host_comm;
 pub mod memory;
 
+// Globals
+#[cfg(feature = "mem_profile")]
+struct ProfiledLockedHeap<const ORDER: usize>(LockedHeap<ORDER>);
+#[cfg(feature = "mem_profile")]
+unsafe impl<const ORDER: usize> alloc::alloc::GlobalAlloc for ProfiledLockedHeap<ORDER> {
+    unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
+        let addr = unsafe { self.0.alloc(layout) };
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") addr as u64);
+        }
+        addr
+    }
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryFree as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") ptr as u64);
+            self.0.dealloc(ptr, layout)
+        }
+    }
+    unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
+        let addr = unsafe { self.0.alloc_zeroed(layout) };
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") addr as u64);
+        }
+        addr
+    }
+    unsafe fn realloc(
+        &self,
+        ptr: *mut u8,
+        layout: core::alloc::Layout,
+        new_size: usize,
+    ) -> *mut u8 {
+        let new_ptr = unsafe { self.0.realloc(ptr, layout, new_size) };
+        unsafe {
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryFree as u16,
+                in("rax") layout.size() as u64,
+                in("rcx") ptr);
+            core::arch::asm!("out dx, al",
+                in("dx") OutBAction::TraceMemoryAlloc as u16,
+                in("rax") new_size as u64,
+                in("rcx") new_ptr);
+        }
+        new_ptr
+    }
+}
+
 // === Globals ===
+#[cfg(not(feature = "mem_profile"))]
 #[global_allocator]
 pub(crate) static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::empty();
+#[cfg(feature = "mem_profile")]
+#[global_allocator]
+pub(crate) static HEAP_ALLOCATOR: ProfiledLockedHeap<32> =
+    ProfiledLockedHeap(LockedHeap::<32>::empty());
 
 pub(crate) static mut GUEST_HANDLE: GuestHandle = GuestHandle::new();
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
@@ -128,7 +190,11 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
 
             let heap_start = (*peb_ptr).guest_heap.ptr as usize;
             let heap_size = (*peb_ptr).guest_heap.size as usize;
-            HEAP_ALLOCATOR
+            #[cfg(not(feature = "mem_profile"))]
+            let heap_allocator = &HEAP_ALLOCATOR;
+            #[cfg(feature = "mem_profile")]
+            let heap_allocator = &HEAP_ALLOCATOR.0;
+            heap_allocator
                 .try_lock()
                 .expect("Failed to access HEAP_ALLOCATOR")
                 .init(heap_start, heap_size);

src/hyperlight_host/Cargo.toml

@@ -30,6 +30,7 @@ libc = { version = "0.2.174" }
 flatbuffers = "25.2.10"
 framehop = { version = "0.13.1", optional = true }
 fallible-iterator = { version = "0.3.0", optional = true }
+blake3 = "1.8.2"
 page_size = "0.6.0"
 termcolor = "1.2.0"
 bitflags = "2.9.1"
@@ -47,7 +48,6 @@ metrics = "0.24.2"
 serde_json = "1.0"
 elfcore = "2.0"
 uuid = { version = "1.17.0", features = ["v4"] }
-blake3 = "1.8.2"
 
 [target.'cfg(windows)'.dependencies]
 windows = { version = "0.61", features = [
@@ -135,6 +135,7 @@ trace_guest = []
 # This feature enables unwinding the guest stack from the host, in
 # order to produce stack traces for debugging or profiling.
 unwind_guest = [ "trace_guest", "dep:framehop", "dep:fallible-iterator", "hyperlight-common/unwind_guest" ]
+mem_profile = [ "unwind_guest", "hyperlight-common/mem_profile" ]
 kvm = ["dep:kvm-bindings", "dep:kvm-ioctls"]
 mshv2 = ["dep:mshv-bindings2", "dep:mshv-ioctls2"]
 mshv3 = ["dep:mshv-bindings3", "dep:mshv-ioctls3"]

src/hyperlight_host/src/sandbox/outb.rs

@@ -257,6 +257,36 @@ fn handle_outb_impl(
                 write_stack(f, &stack);
             })
         }
+        #[cfg(feature = "mem_profile")]
+        OutBAction::TraceMemoryAlloc => {
+            let Ok(stack) = unwind(_hv, mem_mgr.as_ref(), _hv.trace_info_as_ref()) else {
+                return Ok(());
+            };
+            let Ok(amt) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RAX) else {
+                return Ok(());
+            };
+            let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
+                return Ok(());
+            };
+            record_trace_frame(_hv.trace_info_as_ref(), 2u64, |f| {
+                let _ = f.write_all(&ptr.to_ne_bytes());
+                let _ = f.write_all(&amt.to_ne_bytes());
+                write_stack(f, &stack);
+            })
+        }
+        #[cfg(feature = "mem_profile")]
+        OutBAction::TraceMemoryFree => {
+            let Ok(stack) = unwind(_hv, mem_mgr.as_ref(), _hv.trace_info_as_ref()) else {
+                return Ok(());
+            };
+            let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
+                return Ok(());
+            };
+            record_trace_frame(_hv.trace_info_as_ref(), 3u64, |f| {
+                let _ = f.write_all(&ptr.to_ne_bytes());
+                write_stack(f, &stack);
+            })
+        }
     }
 }