Add trace support for recording memory allocations and frees

This allows for producing profiles of memory usage.

Signed-off-by: Lucy Menon <168595099+syntactically@users.noreply.github.com>

Author: syntactically

src/hyperlight_guest/Cargo.toml

@@ -17,6 +17,7 @@ default = ["libc", "printf", "alloca"]
 libc = [] # compile musl libc
 printf = [] # compile printf
 alloca = [] # compile alloca wrapper
+mem_profile = []
 
 [dependencies]
 anyhow = { version = "1.0.45", default-features = false }

src/hyperlight_guest/src/entrypoint.rs

@@ -131,7 +131,11 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_
 
             let heap_start = (*peb_ptr).guestheapData.guestHeapBuffer as usize;
             let heap_size = (*peb_ptr).guestheapData.guestHeapSize 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_guest/src/host_function_call.rs

@@ -37,6 +37,13 @@ pub enum OutBAction {
     Log = 99,
     CallFunction = 101,
     Abort = 102,
+    #[cfg(feature = "mem_profile")]
+    #[allow(dead_code)]
+    TraceRecordStack = 103,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryAlloc = 104,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryFree = 105,
 }
 
 pub fn get_host_value_return_as_void() -> Result<()> {

src/hyperlight_guest/src/lib.rs

@@ -83,8 +83,66 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
 }
 
 // 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 = 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 = 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 = 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
+    }
+}
+#[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());
 
 ///cbindgen:ignore
 #[no_mangle]

src/hyperlight_host/Cargo.toml

@@ -24,6 +24,7 @@ workspace = true
 goblin = { version = "0.9" }
 framehop = { version = "0.13.1", optional = true }
 fallible-iterator = { version = "0.3.0", optional = true }
+blake3 = { version = "1.5.5", optional = true }
 rand = { version = "0.8.5" }
 cfg-if = { version = "1.0.0" }
 libc = { version = "0.2.167" }
@@ -130,6 +131,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" ]
+mem_profile = [ "unwind_guest", "dep:blake3" ]
 kvm = ["dep:kvm-bindings", "dep:kvm-ioctls"]
 mshv = ["dep:mshv-bindings", "dep:mshv-ioctls"]
 inprocess = []

src/hyperlight_host/src/sandbox/outb.rs

@@ -48,6 +48,10 @@ pub(super) enum OutBAction {
     Abort,
     #[cfg(feature = "unwind_guest")]
     TraceRecordStack,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryAlloc,
+    #[cfg(feature = "mem_profile")]
+    TraceMemoryFree,
 }
 
 impl TryFrom<u16> for OutBAction {
@@ -60,6 +64,10 @@ impl TryFrom<u16> for OutBAction {
             102 => Ok(OutBAction::Abort),
             #[cfg(feature = "unwind_guest")]
             103 => Ok(OutBAction::TraceRecordStack),
+            #[cfg(feature = "mem_profile")]
+            104 => Ok(OutBAction::TraceMemoryAlloc),
+            #[cfg(feature = "mem_profile")]
+            105 => Ok(OutBAction::TraceMemoryFree),
             _ => Err(new_error!("Invalid OutB value: {}", val)),
         }
     }
@@ -238,6 +246,36 @@ fn handle_outb_impl(
                 write_stack(f, &stack);
             })
         }
+        #[cfg(feature = "mem_profile")]
+        OutBAction::TraceMemoryAlloc => {
+            let Ok(stack) = unwind(_hv, mem_mgr.as_ref(), &mut _trace_info) 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(&_trace_info, 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(), &mut _trace_info) else {
+                return Ok(());
+            };
+            let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
+                return Ok(());
+            };
+            record_trace_frame(&_trace_info, 3u64, |f| {
+                let _ = f.write_all(&ptr.to_ne_bytes());
+                write_stack(f, &stack);
+            })
+        }
     }
 }