diff --git a/Cargo.lock b/Cargo.lock index 556caa510..f9fe19ee8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -494,6 +494,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dlmalloc" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad5208a115eaba24916f7456929832e310a81518c641f93fee4f89aa93aa3675" +dependencies = [ + "cfg-if", + "libc", + "windows-sys", +] + [[package]] name = "ecsm" version = "0.1.0" @@ -846,6 +857,8 @@ dependencies = [ name = "lambda-vm-syscalls" version = "0.1.0" dependencies = [ + "critical-section", + "dlmalloc", "embedded-alloc", "getrandom 0.2.16", "getrandom 0.3.4", diff --git a/syscalls/Cargo.lock b/syscalls/Cargo.lock index 34e481dd8..1921ff54a 100644 --- a/syscalls/Cargo.lock +++ b/syscalls/Cargo.lock @@ -64,6 +64,17 @@ dependencies = [ "crypto-common", ] +[[package]] +name = "dlmalloc" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad5208a115eaba24916f7456929832e310a81518c641f93fee4f89aa93aa3675" +dependencies = [ + "cfg-if", + "libc", + "windows-sys", +] + [[package]] name = "embedded-alloc" version = "0.6.0" @@ -128,6 +139,8 @@ dependencies = [ name = "lambda-vm-syscalls" version = "0.1.0" dependencies = [ + "critical-section", + "dlmalloc", "embedded-alloc", "getrandom 0.2.17", "getrandom 0.3.4", @@ -379,6 +392,21 @@ dependencies = [ "wit-bindgen", ] +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + [[package]] name = "wit-bindgen" version = "0.57.1" diff --git a/syscalls/Cargo.toml b/syscalls/Cargo.toml index 0460a2435..f1601ee35 100644 --- a/syscalls/Cargo.toml +++ b/syscalls/Cargo.toml @@ -4,6 +4,10 @@ version = "0.1.0" edition = "2024" [dependencies] +# embedded-alloc TLSF heap: the previous default guest allocator, now selectable via +# the `tlsf-alloc` feature. Kept a hard dep (like the codebase's other allocator +# toggles) so switching needs no dependency edits; it's only linked into the guest and +# dropped when dlmalloc (the default) is selected. embedded-alloc = "0.6" riscv = { version = "0.15", features = ["critical-section-single-hart"] } thiserror = "1.0" @@ -11,6 +15,17 @@ getrandom = { version = "0.3.4", default-features = false } getrandom_v2 = {version = "0.2.15", features = ["custom"], package = "getrandom"} lazy_static = "1.5.0" rand = "0.9.2" +# Default guest allocator: Doug Lea's malloc. Cheaper per alloc/free than TLSF on zkVM +# guests (fewer guest cycles) while still reclaiming freed memory (no OOM). `critical- +# section` gives the Sync a #[global_allocator] static needs; its single-hart impl comes +# from `riscv` above (same as embedded-alloc's TLSF heap uses). See `src/allocator.rs`. +dlmalloc = { version = "0.2.14", default-features = false } +critical-section = "1.2" + +[features] +# Fallback/BENCH: use the embedded-alloc TLSF heap instead of the default dlmalloc guest +# allocator (to A/B the two or fall back). Off by default. See `src/allocator.rs`. +tlsf-alloc = [] [dev-dependencies] keccak = "0.1" diff --git a/syscalls/src/allocator.rs b/syscalls/src/allocator.rs index 78b2933e5..c99482b45 100644 --- a/syscalls/src/allocator.rs +++ b/syscalls/src/allocator.rs @@ -1,23 +1,164 @@ -use embedded_alloc::TlsfHeap as Heap; use riscv as _; -// Only the guest routes Rust allocations through this heap; on host (e.g. -// `cargo test` for the sponge's differential tests) the attribute would hijack -// the test harness's allocator with a never-initialized heap and abort. -#[cfg_attr(target_arch = "riscv64", global_allocator)] -static HEAP: Heap = Heap::empty(); - const MAX_MEMORY_SIZE: usize = 0xC000_0000; const WORD_SIZE: usize = 4; -pub fn init_allocator() { - { - unsafe extern "C" { - static _end: u8; +// Guest global allocator, selectable at build time: +// - default: Doug Lea's malloc (dlmalloc), backed by a bump "system" provider that +// hands dlmalloc page-aligned chunks from the guest heap. dlmalloc does all +// sub-allocation churn itself, so — unlike a raw bump allocator — it reuses freed +// memory (no OOM) while executing fewer guest instructions per alloc/free than TLSF +// (Bencik's ZisK allocator comparison; measured cheaper on our ethrex proving too). +// - `tlsf-alloc` feature: embedded-alloc's TLSF heap — the previous default, kept one +// flag away for A/B comparison or fallback. +// +// Only the guest installs a #[global_allocator]; on host (e.g. `cargo test` for the +// sponge's differential tests) the attribute would hijack the test harness's +// allocator with a never-initialized heap and abort. + +#[cfg(not(feature = "tlsf-alloc"))] +mod imp { + use core::alloc::{GlobalAlloc, Layout}; + use core::cell::RefCell; + use core::sync::atomic::{AtomicUsize, Ordering}; + use critical_section::Mutex; + use dlmalloc::{Allocator, Dlmalloc}; + + // Page granularity dlmalloc requests memory in. Must be a power of two; the guest + // heap region is 3 GiB so the value only affects the segment rounding below. + const PAGE_SIZE: usize = 4096; + + // The "system" side of dlmalloc: instead of mmap/sbrk (absent on the guest) it + // bump-allocates page-aligned segments from the single contiguous heap region + // [_end, MAX_MEMORY_SIZE). It never releases a segment (`free`/`free_part`/ + // `remap` all decline) — dlmalloc itself owns all reuse of freed *user* + // allocations against this fixed backing store, which is what keeps churny + // workloads OOM-free unlike a raw bump allocator. + struct BumpSystem; + + // Single-hart guest → `Relaxed` atomics are contention-free. + static HEAP_POS: AtomicUsize = AtomicUsize::new(0); + static HEAP_END: AtomicUsize = AtomicUsize::new(0); + + unsafe impl Allocator for BumpSystem { + fn alloc(&self, size: usize) -> (*mut u8, usize, u32) { + // Round up to a page so consecutive segments stay page-aligned. + let size = size.wrapping_add(PAGE_SIZE - 1) & !(PAGE_SIZE - 1); + let pos = HEAP_POS.load(Ordering::Relaxed); + match pos.checked_add(size) { + Some(new_pos) if new_pos <= HEAP_END.load(Ordering::Relaxed) => { + HEAP_POS.store(new_pos, Ordering::Relaxed); + // flags = 0: a plain external segment (never partially released). + (pos as *mut u8, size, 0) + } + // Out of heap → null makes dlmalloc return null → handle_alloc_error. + _ => (core::ptr::null_mut(), 0, 0), + } + } + + fn remap(&self, _ptr: *mut u8, _old: usize, _new: usize, _can_move: bool) -> *mut u8 { + core::ptr::null_mut() + } + + fn free_part(&self, _ptr: *mut u8, _old: usize, _new: usize) -> bool { + false + } + + fn free(&self, _ptr: *mut u8, _size: usize) -> bool { + false + } + + fn can_release_part(&self, _flags: u32) -> bool { + false + } + + fn allocates_zeros(&self) -> bool { + // Guest memory is zero-initialized and this provider never reuses a + // segment, so system-fresh bytes read as 0 → dlmalloc's calloc skips the + // memset for system-fresh memory (it still zeroes recycled blocks itself). + true + } + + fn page_size(&self) -> usize { + PAGE_SIZE + } + } + + // Dlmalloc is Send but !Sync, so it can't sit in a static directly. A single-hart + // critical section serializes access and supplies the Sync a #[global_allocator] + // static requires — the same primitive embedded-alloc's TLSF heap uses. + static DLMALLOC: Mutex>> = + Mutex::new(RefCell::new(Dlmalloc::new_with_allocator(BumpSystem))); + + struct DlGlobal; + + #[cfg_attr(target_arch = "riscv64", global_allocator)] + static ALLOC: DlGlobal = DlGlobal; + + pub fn init(heap_start: usize, heap_end: usize) { + HEAP_POS.store(heap_start, Ordering::Relaxed); + HEAP_END.store(heap_end, Ordering::Relaxed); + } + + unsafe impl GlobalAlloc for DlGlobal { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .malloc(layout.size(), layout.align()) + }) } - let heap_pos: usize = unsafe { (&_end) as *const u8 as usize }; - unsafe { HEAP.init(heap_pos, MAX_MEMORY_SIZE - heap_pos) } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .free(ptr, layout.size(), layout.align()) + }) + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC + .borrow(cs) + .borrow_mut() + .calloc(layout.size(), layout.align()) + }) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + critical_section::with(|cs| unsafe { + DLMALLOC.borrow(cs).borrow_mut().realloc( + ptr, + layout.size(), + layout.align(), + new_size, + ) + }) + } + } +} + +#[cfg(feature = "tlsf-alloc")] +mod imp { + use embedded_alloc::TlsfHeap as Heap; + + #[cfg_attr(target_arch = "riscv64", global_allocator)] + static HEAP: Heap = Heap::empty(); + + pub fn init(heap_start: usize, heap_end: usize) { + unsafe { HEAP.init(heap_start, heap_end - heap_start) } + } +} + +pub fn init_allocator() { + unsafe extern "C" { + static _end: u8; } + let heap_pos: usize = unsafe { (&_end) as *const u8 as usize }; + imp::init(heap_pos, MAX_MEMORY_SIZE); } /// # Safety @@ -26,8 +167,8 @@ pub fn init_allocator() { /// It is only for rust std internal uses #[unsafe(no_mangle)] pub unsafe extern "C" fn sys_alloc_aligned(bytes: usize, align: usize) -> *mut u8 { - use core::alloc::GlobalAlloc; - unsafe { HEAP.alloc(core::alloc::Layout::from_size_align(bytes, align).unwrap()) } + // Route through whichever `#[global_allocator]` is installed (dlmalloc or TLSF). + unsafe { std::alloc::alloc(core::alloc::Layout::from_size_align(bytes, align).unwrap()) } } /// # Safety