feat(guest): use mimalloc global allocator

Author: andreiltd

.gitmodules

@@ -2,3 +2,6 @@
 	path = src/hyperlight_guest_bin/third_party/picolibc
 	url = https://github.com/picolibc/picolibc.git
 	shallow = true
+[submodule "src/hyperlight_guest_bin/third_party/mimalloc"]
+	path = src/hyperlight_guest_bin/third_party/mimalloc
+	url = https://github.com/microsoft/mimalloc

Cargo.lock

@@ -27,6 +27,7 @@ hyperlight-guest-tracing = { workspace = true, default-features = false }
 buddy_system_allocator = "0.11.0"
 log = { version = "0.4", default-features = false }
 spin = "0.10.0"
+libc_alloc = "1.0.7"
 
 [lints]
 workspace = true

src/hyperlight_guest_bin/Cargo.toml

@@ -209,6 +209,7 @@ fn cc_build(picolibc_dir: &PathBuf, target: &str) -> Result<cc::Build> {
         .flag("-fno-common")
         .flag("-fno-builtin")
         .flag("-Wno-unused-command-line-argument")
+        .flag("-Wno-unused-parameter")
         .flag("-fdiagnostics-color=always")
         .flag("-Wall")
         .flag("-Winvalid-pch")
@@ -265,6 +266,16 @@ fn cc_build(picolibc_dir: &PathBuf, target: &str) -> Result<cc::Build> {
     Ok(build)
 }
 
+fn add_mimalloc(build: &mut cc::Build, mimalloc_dir: &Path, _target: &str) -> Result<()> {
+    build.define("MI_MALLOC_OVERRIDE", "1");
+    build.define("__wasi__", "1");
+    build.include(mimalloc_dir.join("include"));
+    build.include(mimalloc_dir.join("src"));
+    build.file(mimalloc_dir.join("src/static.c"));
+
+    Ok(())
+}
+
 fn add_libc(build: &mut cc::Build, picolibc_dir: &Path, target: &str) -> Result<()> {
     let base = LIBC_FILES.iter();
     let files = match target {
@@ -277,7 +288,7 @@ fn add_libc(build: &mut cc::Build, picolibc_dir: &Path, target: &str) -> Result<
         build.file(&source_path);
     }
 
-    build.file("c/clock.c");
+    build.files(&["c/clock.c", "c/mimalloc.c"]);
     Ok(())
 }
 
@@ -331,9 +342,11 @@ fn cargo_main() -> Result<()> {
         }
 
         let picolibc_dir = PathBuf::from("third_party/picolibc");
+        let mimalloc_dir = PathBuf::from("third_party/mimalloc");
         let mut build = cc_build(&picolibc_dir, &target)?;
 
         add_libc(&mut build, &picolibc_dir, &target)?;
+        add_mimalloc(&mut build, &mimalloc_dir, &target)?;
 
         if want_libm {
             add_libm(&mut build, &picolibc_dir, &target)?;

src/hyperlight_guest_bin/build.rs

@@ -434,7 +434,6 @@ pub(crate) const LIBC_FILES: &[&str] = &[
     "ubsan/ubsan_val_to_string.c",
     "ubsan/ubsan_val_to_umax.c",
     "ubsan/ubsan_warning.c",
-    "picolib/picosbrk.c",
     "picolib/dso_handle.c",
     "picolib/getauxval.c",
     "picolib/machine/x86/interrupt.c",

src/hyperlight_guest_bin/build_files.rs

@@ -0,0 +1,85 @@
+#include <stdbool.h>
+#include <stdlib.h>
+#include <stdio.h>
+#include <unistd.h>
+#include <errno.h>
+#include <mimalloc.h>
+
+// -----------------------------------------------------------------------------
+// Static variables for sbrk region_
+// -----------------------------------------------------------------------------
+static char *heap_start = NULL;
+static char *heap_end   = NULL;
+static char *brk_ptr    = NULL;
+
+/*
+ * Initialize a mimalloc arena on the donated memory region.
+ * After this call, all allocations use only the pre‐reserved arena.
+ */
+void init_arena(void *donated_ptr, size_t donated_size) {
+    /* do not use OS memory for allocation */
+    mi_option_set_enabled(mi_option_limit_os_alloc, true);
+
+    mi_arena_id_t arena_id = 0;
+    bool ok = mi_manage_os_memory_ex(
+        donated_ptr,
+        donated_size,
+        1 /* committed */,
+        0 /* large     */,
+        0 /* zero      */,
+        -1 /* numa_node */,
+        1 /* exclusive */,
+        &arena_id
+    );
+
+    if (!ok) {
+        fprintf(stderr, "mi_manage_os_memory_ex failed\n");
+        abort();
+    }
+
+    /* create a heap in the donated arena and set it as default */
+    mi_heap_t *heap = mi_heap_new_in_arena(arena_id);
+    mi_heap_set_default(heap);
+}
+
+/*
+ * Register the heap region to be used by sbrk().
+ * Must be called before any sbrk() invocation.
+ */
+void init_sbrk(void *start, size_t size) {
+    heap_start = (char*)start;
+    heap_end   = heap_start + size;
+    brk_ptr    = heap_start;
+}
+
+/*
+ * Unix‐style sbrk implementation over the registered heap region.
+ * On success returns the previous break; on failure returns (void*)-1
+ * and sets errno = ENOMEM.
+ */
+void *sbrk(ptrdiff_t incr) {
+    if (heap_start == NULL) {
+        errno = ENOMEM;
+        return (void*)-1;
+    }
+
+    if (incr < 0) {
+        /* shrink: ensure we don't move below heap_start */
+        size_t dec = (size_t)(-incr);
+        if ((size_t)(brk_ptr - heap_start) < dec) {
+            errno = ENOMEM;
+            return (void*)-1;
+        }
+    } else {
+        /* grow: ensure we don't move past heap_end */
+        size_t inc = (size_t)incr;
+        if ((size_t)(heap_end - brk_ptr) < inc) {
+            errno = ENOMEM;
+            return (void*)-1;
+        }
+    }
+
+    void *old_brk = brk_ptr;
+    brk_ptr += incr;
+    return old_brk;
+}

src/hyperlight_guest_bin/c/mimalloc.c

@@ -18,9 +18,9 @@ limitations under the License.
 // === Dependencies ===
 extern crate alloc;
 
+use core::ffi::c_void;
 use alloc::string::ToString;
 
-use buddy_system_allocator::LockedHeap;
 #[cfg(target_arch = "x86_64")]
 use exceptions::{gdt::load_gdt, idtr::load_idt};
 use guest_function::call::dispatch_function;
@@ -34,6 +34,7 @@ use hyperlight_guest::exit::{abort_with_code_and_message, halt};
 use hyperlight_guest::guest_handle::handle::GuestHandle;
 use hyperlight_guest_tracing::{trace, trace_function};
 use log::LevelFilter;
+use libc_alloc::LibcAlloc;
 use spin::Once;
 
 // === Modules ===
@@ -54,75 +55,17 @@ pub mod guest_function {
 
 pub mod guest_logger;
 pub mod host_comm;
-pub mod memory;
 pub mod paging;
 
 #[cfg(feature = "libc")]
 mod host_bridge;
 
-// 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());
+static ALLOCATOR: LibcAlloc = LibcAlloc;
 
 pub(crate) static mut GUEST_HANDLE: GuestHandle = GuestHandle::new();
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister =
@@ -154,6 +97,8 @@ fn panic(info: &core::panic::PanicInfo) -> ! {
 unsafe extern "C" {
     fn hyperlight_main();
     fn srand(seed: u32);
+    fn init_sbrk(donated_ptr: *const c_void, donated_size: usize);
+    fn init_arena(donated_ptr: *const c_void, donated_size: usize);
 }
 
 static INIT: Once = Once::new();
@@ -188,16 +133,10 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
                 load_idt();
             }
 
-            let heap_start = (*peb_ptr).guest_heap.ptr as usize;
+            let heap_start = (*peb_ptr).guest_heap.ptr as *const c_void;
             let heap_size = (*peb_ptr).guest_heap.size as usize;
-            #[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);
+            init_arena(heap_start, heap_size);
+            init_sbrk(heap_start, heap_size);
 
             OS_PAGE_SIZE = ops as u32;