Use talc as the guest heap allocator

CHANGELOG.md

@@ -31,6 +31,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 * A new sandbox writes only the guest code, PEB and init data into its memory.
   The host populates other pages on first guest touch, which speeds up sandbox
   creation.
+* Use `talc` as the guest heap allocator. Guest init writes only the start and
+  end of the heap, which speeds up sandbox creation.
 
 * **Breaking:** The transport arena sits one unmapped guard page below the reserved scratch-top pages, at an address independent of the scratch size, and the page tables at the scratch base. The guest's page allocator stops at the arena GPA the host publishes in the scratch-top metadata, so guests must be rebuilt. `hyperlight_common::layout::scratch_allocator_limit_gpa` is renamed `scratch_reserved_top_gpa`, the start of the reserved pages. An allocator that stops there runs into the arena. The minimum scratch size grows by that guard page. Snapshot ABI 7: existing snapshots must be regenerated.
 

src/hyperlight_guest_bin/Cargo.toml

@@ -25,7 +25,7 @@ hyperlight-common = { workspace = true, default-features = false }
 hyperlight-guest-tracing = { workspace = true, default-features = false }
 hyperlight-guest-macro = { workspace = true, default-features = false, optional = true }
 hyperlight-libc = { workspace = true, default-features = false, optional = true }
-buddy_system_allocator = "0.13.0"
+talc = { version = "5.1.1", default-features = false }
 log = { version = "0.4", default-features = false }
 linkme = { version = "0.3.36", optional = true }
 spin = "0.12.0"

src/hyperlight_guest_bin/src/lib.rs

@@ -8,7 +8,6 @@ extern crate alloc;
 use core::fmt::Write;
 
 use arch::dispatch::dispatch_function;
-use buddy_system_allocator::LockedHeap;
 use guest_function::register::GuestFunctionRegister;
 use guest_logger::init_logger;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
@@ -18,6 +17,8 @@ use hyperlight_common::mem::HyperlightPEB;
 use hyperlight_common::outb::OutBAction;
 use hyperlight_guest::exit::write_abort;
 use hyperlight_guest::guest_handle::handle::GuestHandle;
+use talc::TalcLock;
+use talc::source::Manual;
 
 // === Modules ===
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/mod.rs")]
@@ -50,10 +51,12 @@ mod init;
 pub use hyperlight_libc as libc;
 
 // Globals
+type GuestHeap = TalcLock<spin::Mutex<()>, Manual>;
+
 #[cfg(all(feature = "mem_profile", target_arch = "x86_64"))]
-struct ProfiledLockedHeap<const ORDER: usize>(LockedHeap<ORDER>);
+struct ProfiledLockedHeap(GuestHeap);
 #[cfg(all(feature = "mem_profile", target_arch = "x86_64"))]
-unsafe impl<const ORDER: usize> alloc::alloc::GlobalAlloc for ProfiledLockedHeap<ORDER> {
+unsafe impl alloc::alloc::GlobalAlloc for ProfiledLockedHeap {
     unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
         let addr = unsafe { self.0.alloc(layout) };
         unsafe {
@@ -107,11 +110,10 @@ unsafe impl<const ORDER: usize> alloc::alloc::GlobalAlloc for ProfiledLockedHeap
 // === Globals ===
 #[cfg(not(all(feature = "mem_profile", target_arch = "x86_64")))]
 #[global_allocator]
-pub(crate) static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::empty();
+pub(crate) static HEAP_ALLOCATOR: GuestHeap = GuestHeap::new(Manual);
 #[cfg(all(feature = "mem_profile", target_arch = "x86_64"))]
 #[global_allocator]
-pub(crate) static HEAP_ALLOCATOR: ProfiledLockedHeap<32> =
-    ProfiledLockedHeap(LockedHeap::<32>::empty());
+pub(crate) static HEAP_ALLOCATOR: ProfiledLockedHeap = ProfiledLockedHeap(GuestHeap::new(Manual));
 
 pub static mut GUEST_HANDLE: GuestHandle = GuestHandle::new();
 pub(crate) static mut REGISTERED_GUEST_FUNCTIONS: GuestFunctionRegister<GuestFunc> =
@@ -268,10 +270,13 @@ pub(crate) extern "C" fn generic_init(
         let heap_allocator = &HEAP_ALLOCATOR;
         #[cfg(all(feature = "mem_profile", target_arch = "x86_64"))]
         let heap_allocator = &HEAP_ALLOCATOR.0;
+        // SAFETY: the host reserves the PEB heap range for the guest heap, and
+        // init claims it once. Only this allocator writes to it.
         heap_allocator
             .try_lock()
             .expect("Failed to access HEAP_ALLOCATOR")
-            .init(heap_start, heap_size);
+            .claim(heap_start as *mut u8, heap_size)
+            .expect("Failed to claim guest heap");
         peb_ptr
     };