From b51db53ffa204d6666787bae3149b75b85c919cd Mon Sep 17 00:00:00 2001 From: Lucy Menon <168595099+syntactically@users.noreply.github.com> Date: Fri, 6 Dec 2024 12:41:55 +0000 Subject: [PATCH] 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> --- src/hyperlight_guest/Cargo.toml | 1 + src/hyperlight_guest/src/entrypoint.rs | 6 +- .../src/host_function_call.rs | 7 +++ src/hyperlight_guest/src/lib.rs | 58 +++++++++++++++++++ src/hyperlight_host/Cargo.toml | 2 + src/hyperlight_host/src/sandbox/outb.rs | 38 ++++++++++++ 6 files changed, 111 insertions(+), 1 deletion(-) diff --git a/src/hyperlight_guest/Cargo.toml b/src/hyperlight_guest/Cargo.toml index 8f218b1f7..33346a71d 100644 --- a/src/hyperlight_guest/Cargo.toml +++ b/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 } diff --git a/src/hyperlight_guest/src/entrypoint.rs b/src/hyperlight_guest/src/entrypoint.rs index f6fc15940..889f36b10 100644 --- a/src/hyperlight_guest/src/entrypoint.rs +++ b/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); diff --git a/src/hyperlight_guest/src/host_function_call.rs b/src/hyperlight_guest/src/host_function_call.rs index f5e96a7ba..6d19f8ac0 100644 --- a/src/hyperlight_guest/src/host_function_call.rs +++ b/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<()> { diff --git a/src/hyperlight_guest/src/lib.rs b/src/hyperlight_guest/src/lib.rs index 5577e4c41..9e76cee01 100644 --- a/src/hyperlight_guest/src/lib.rs +++ b/src/hyperlight_guest/src/lib.rs @@ -83,8 +83,66 @@ fn panic(info: &core::panic::PanicInfo) -> ! { } // Globals +#[cfg(feature = "mem_profile")] +struct ProfiledLockedHeap(LockedHeap); +#[cfg(feature = "mem_profile")] +unsafe impl alloc::alloc::GlobalAlloc for ProfiledLockedHeap { + 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] diff --git a/src/hyperlight_host/Cargo.toml b/src/hyperlight_host/Cargo.toml index e357ad295..702b8f766 100644 --- a/src/hyperlight_host/Cargo.toml +++ b/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 = [] diff --git a/src/hyperlight_host/src/sandbox/outb.rs b/src/hyperlight_host/src/sandbox/outb.rs index 68066758a..e30419b34 100644 --- a/src/hyperlight_host/src/sandbox/outb.rs +++ b/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 for OutBAction { @@ -60,6 +64,10 @@ impl TryFrom 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); + }) + } } }