Skip to content

Commit b51db53

Browse files
committed
Add trace support for recording memory allocations and frees
This allows for producing profiles of memory usage. Signed-off-by: Lucy Menon <[email protected]>
1 parent d241502 commit b51db53

File tree

6 files changed

+111
-1
lines changed

6 files changed

+111
-1
lines changed

src/hyperlight_guest/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ default = ["libc", "printf", "alloca"]
1717
libc = [] # compile musl libc
1818
printf = [] # compile printf
1919
alloca = [] # compile alloca wrapper
20+
mem_profile = []
2021

2122
[dependencies]
2223
anyhow = { version = "1.0.45", default-features = false }

src/hyperlight_guest/src/entrypoint.rs

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,11 @@ pub extern "win64" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_
131131

132132
let heap_start = (*peb_ptr).guestheapData.guestHeapBuffer as usize;
133133
let heap_size = (*peb_ptr).guestheapData.guestHeapSize as usize;
134-
HEAP_ALLOCATOR
134+
#[cfg(not(feature = "mem_profile"))]
135+
let heap_allocator = &HEAP_ALLOCATOR;
136+
#[cfg(feature = "mem_profile")]
137+
let heap_allocator = &HEAP_ALLOCATOR.0;
138+
heap_allocator
135139
.try_lock()
136140
.expect("Failed to access HEAP_ALLOCATOR")
137141
.init(heap_start, heap_size);

src/hyperlight_guest/src/host_function_call.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,13 @@ pub enum OutBAction {
3737
Log = 99,
3838
CallFunction = 101,
3939
Abort = 102,
40+
#[cfg(feature = "mem_profile")]
41+
#[allow(dead_code)]
42+
TraceRecordStack = 103,
43+
#[cfg(feature = "mem_profile")]
44+
TraceMemoryAlloc = 104,
45+
#[cfg(feature = "mem_profile")]
46+
TraceMemoryFree = 105,
4047
}
4148

4249
pub fn get_host_value_return_as_void() -> Result<()> {

src/hyperlight_guest/src/lib.rs

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,66 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
8383
}
8484

8585
// Globals
86+
#[cfg(feature = "mem_profile")]
87+
struct ProfiledLockedHeap<const ORDER: usize>(LockedHeap<ORDER>);
88+
#[cfg(feature = "mem_profile")]
89+
unsafe impl<const ORDER: usize> alloc::alloc::GlobalAlloc for ProfiledLockedHeap<ORDER> {
90+
unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
91+
let addr = self.0.alloc(layout);
92+
unsafe {
93+
core::arch::asm!("out dx, al",
94+
in("dx") OutBAction::TraceMemoryAlloc as u16,
95+
in("rax") layout.size() as u64,
96+
in("rcx") addr as u64);
97+
}
98+
addr
99+
}
100+
unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
101+
unsafe {
102+
core::arch::asm!("out dx, al",
103+
in("dx") OutBAction::TraceMemoryFree as u16,
104+
in("rax") layout.size() as u64,
105+
in("rcx") ptr as u64);
106+
}
107+
self.0.dealloc(ptr, layout)
108+
}
109+
unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
110+
let addr = self.0.alloc_zeroed(layout);
111+
unsafe {
112+
core::arch::asm!("out dx, al",
113+
in("dx") OutBAction::TraceMemoryAlloc as u16,
114+
in("rax") layout.size() as u64,
115+
in("rcx") addr as u64);
116+
}
117+
addr
118+
}
119+
unsafe fn realloc(
120+
&self,
121+
ptr: *mut u8,
122+
layout: core::alloc::Layout,
123+
new_size: usize,
124+
) -> *mut u8 {
125+
let new_ptr = self.0.realloc(ptr, layout, new_size);
126+
unsafe {
127+
core::arch::asm!("out dx, al",
128+
in("dx") OutBAction::TraceMemoryFree as u16,
129+
in("rax") layout.size() as u64,
130+
in("rcx") ptr);
131+
core::arch::asm!("out dx, al",
132+
in("dx") OutBAction::TraceMemoryAlloc as u16,
133+
in("rax") new_size as u64,
134+
in("rcx") new_ptr);
135+
}
136+
new_ptr
137+
}
138+
}
139+
#[cfg(not(feature = "mem_profile"))]
86140
#[global_allocator]
87141
pub(crate) static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::empty();
142+
#[cfg(feature = "mem_profile")]
143+
#[global_allocator]
144+
pub(crate) static HEAP_ALLOCATOR: ProfiledLockedHeap<32> =
145+
ProfiledLockedHeap(LockedHeap::<32>::empty());
88146

89147
///cbindgen:ignore
90148
#[no_mangle]

src/hyperlight_host/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ workspace = true
2424
goblin = { version = "0.9" }
2525
framehop = { version = "0.13.1", optional = true }
2626
fallible-iterator = { version = "0.3.0", optional = true }
27+
blake3 = { version = "1.5.5", optional = true }
2728
rand = { version = "0.8.5" }
2829
cfg-if = { version = "1.0.0" }
2930
libc = { version = "0.2.167" }
@@ -130,6 +131,7 @@ trace_guest = []
130131
# This feature enables unwinding the guest stack from the host, in
131132
# order to produce stack traces for debugging or profiling.
132133
unwind_guest = [ "trace_guest", "dep:framehop", "dep:fallible-iterator" ]
134+
mem_profile = [ "unwind_guest", "dep:blake3" ]
133135
kvm = ["dep:kvm-bindings", "dep:kvm-ioctls"]
134136
mshv = ["dep:mshv-bindings", "dep:mshv-ioctls"]
135137
inprocess = []

src/hyperlight_host/src/sandbox/outb.rs

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,10 @@ pub(super) enum OutBAction {
4848
Abort,
4949
#[cfg(feature = "unwind_guest")]
5050
TraceRecordStack,
51+
#[cfg(feature = "mem_profile")]
52+
TraceMemoryAlloc,
53+
#[cfg(feature = "mem_profile")]
54+
TraceMemoryFree,
5155
}
5256

5357
impl TryFrom<u16> for OutBAction {
@@ -60,6 +64,10 @@ impl TryFrom<u16> for OutBAction {
6064
102 => Ok(OutBAction::Abort),
6165
#[cfg(feature = "unwind_guest")]
6266
103 => Ok(OutBAction::TraceRecordStack),
67+
#[cfg(feature = "mem_profile")]
68+
104 => Ok(OutBAction::TraceMemoryAlloc),
69+
#[cfg(feature = "mem_profile")]
70+
105 => Ok(OutBAction::TraceMemoryFree),
6371
_ => Err(new_error!("Invalid OutB value: {}", val)),
6472
}
6573
}
@@ -238,6 +246,36 @@ fn handle_outb_impl(
238246
write_stack(f, &stack);
239247
})
240248
}
249+
#[cfg(feature = "mem_profile")]
250+
OutBAction::TraceMemoryAlloc => {
251+
let Ok(stack) = unwind(_hv, mem_mgr.as_ref(), &mut _trace_info) else {
252+
return Ok(());
253+
};
254+
let Ok(amt) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RAX) else {
255+
return Ok(());
256+
};
257+
let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
258+
return Ok(());
259+
};
260+
record_trace_frame(&_trace_info, 2u64, |f| {
261+
let _ = f.write_all(&ptr.to_ne_bytes());
262+
let _ = f.write_all(&amt.to_ne_bytes());
263+
write_stack(f, &stack);
264+
})
265+
}
266+
#[cfg(feature = "mem_profile")]
267+
OutBAction::TraceMemoryFree => {
268+
let Ok(stack) = unwind(_hv, mem_mgr.as_ref(), &mut _trace_info) else {
269+
return Ok(());
270+
};
271+
let Ok(ptr) = _hv.read_trace_reg(crate::hypervisor::TraceRegister::RCX) else {
272+
return Ok(());
273+
};
274+
record_trace_frame(&_trace_info, 3u64, |f| {
275+
let _ = f.write_all(&ptr.to_ne_bytes());
276+
write_stack(f, &stack);
277+
})
278+
}
241279
}
242280
}
243281

0 commit comments

Comments
 (0)