From dc1aaced93d7fe962dfd513aec5ddc24ae0f3e8c Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sat, 8 Aug 2026 11:31:23 +0800 Subject: [PATCH 1/4] feat(fspy): add fspy_alloc, a lock-free async-signal-safe allocator for the preload library MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The preload library interposes libc functions that POSIX declares async-signal-safe (open, stat, execve, ...), so its Rust allocations must not take libc malloc's locks: a signal handler, or the child of fork() in a multithreaded process, would deadlock on locks held by suspended or vanished threads. fspy_alloc routes every allocation in the preload cdylib through a lock-free size-class pool: power-of-two classes carve blocks from 1 MiB slabs, freed blocks recycle through per-class Treiber free lists made ABA-resistant by a 40-bit generation tag, and larger or over-aligned requests map directly. All memory comes from anonymous mappings issued as raw syscalls via rustix's linux_raw backend — on Linux the allocator relies on nothing from libc, discovering even the page size with an mprotect probe (no getauxval, no /proc, no minimum kernel version). macOS goes through the thin libSystem stubs, its only syscall interface. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 8 + Cargo.toml | 2 + crates/fspy_alloc/Cargo.toml | 14 + crates/fspy_alloc/src/class.rs | 38 ++ crates/fspy_alloc/src/lib.rs | 124 +++++ crates/fspy_alloc/src/mapping.rs | 62 +++ crates/fspy_alloc/src/mmap.rs | 211 +++++++ crates/fspy_alloc/src/pool.rs | 705 ++++++++++++++++++++++++ crates/fspy_alloc/src/slab.rs | 203 +++++++ crates/fspy_alloc/src/sys.rs | 25 + crates/fspy_alloc/tests/global_alloc.rs | 84 +++ crates/fspy_preload_unix/Cargo.toml | 1 + crates/fspy_preload_unix/src/lib.rs | 10 + 13 files changed, 1487 insertions(+) create mode 100644 crates/fspy_alloc/Cargo.toml create mode 100644 crates/fspy_alloc/src/class.rs create mode 100644 crates/fspy_alloc/src/lib.rs create mode 100644 crates/fspy_alloc/src/mapping.rs create mode 100644 crates/fspy_alloc/src/mmap.rs create mode 100644 crates/fspy_alloc/src/pool.rs create mode 100644 crates/fspy_alloc/src/slab.rs create mode 100644 crates/fspy_alloc/src/sys.rs create mode 100644 crates/fspy_alloc/tests/global_alloc.rs diff --git a/Cargo.lock b/Cargo.lock index 00ee74ab3..6182d52d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1243,6 +1243,13 @@ dependencies = [ "winsafe 0.0.27", ] +[[package]] +name = "fspy_alloc" +version = "0.0.0" +dependencies = [ + "rustix", +] + [[package]] name = "fspy_benchmark" version = "0.0.0" @@ -1299,6 +1306,7 @@ dependencies = [ "artifact_profile", "bstr", "ctor", + "fspy_alloc", "fspy_shared", "fspy_shared_unix", "libc", diff --git a/Cargo.toml b/Cargo.toml index ee6abee66..4385383c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -73,6 +73,7 @@ materialized_artifact = { path = "crates/materialized_artifact" } materialized_artifact_build = { path = "crates/materialized_artifact_build" } flate2 = "1.0.35" fspy = { path = "crates/fspy" } +fspy_alloc = { path = "crates/fspy_alloc" } fspy_benchmark_launcher = { path = "crates/fspy_benchmark_launcher", artifact = "bin" } fspy_benchmark_target = { path = "crates/fspy_benchmark_target", artifact = "bin" } fspy_detours_sys = { path = "crates/fspy_detours_sys" } @@ -120,6 +121,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" +rustix = { version = "1", default-features = false, features = ["mm", "param"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml new file mode 100644 index 000000000..809d0f7e8 --- /dev/null +++ b/crates/fspy_alloc/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "fspy_alloc" +edition = "2024" +license.workspace = true +publish = false + +[lib] +doctest = false + +[target.'cfg(unix)'.dependencies] +rustix = { workspace = true } + +[lints] +workspace = true diff --git a/crates/fspy_alloc/src/class.rs b/crates/fspy_alloc/src/class.rs new file mode 100644 index 000000000..df8d3797a --- /dev/null +++ b/crates/fspy_alloc/src/class.rs @@ -0,0 +1,38 @@ +//! Size-class policy: which layouts the pool serves, and at what block size. + +use core::alloc::Layout; + +const MIN_CLASS_SHIFT: u32 = 4; +const MAX_CLASS_SHIFT: u32 = 16; +pub const CLASS_COUNT: usize = (MAX_CLASS_SHIFT - MIN_CLASS_SHIFT) as usize + 1; +const MIN_BLOCK_SIZE: usize = 1 << MIN_CLASS_SHIFT; +pub const MAX_BLOCK_SIZE: usize = 1 << MAX_CLASS_SHIFT; +/// Block areas start at this alignment within a slab, making it the largest +/// alignment the pool can serve; stricter layouts map directly. +pub const MAX_POOL_ALIGN: usize = 4096; + +pub const fn block_size(class: usize) -> usize { + 1 << (MIN_CLASS_SHIFT as usize + class) +} + +/// Returns the size class for `layout`, or `None` if the request must be +/// mapped directly (too large or over-aligned). +pub const fn class_of(layout: Layout) -> Option { + if layout.align() > MAX_POOL_ALIGN { + return None; + } + let mut size = layout.size(); + // A block of `size >= align` at a `min(block size, 4 KiB)` boundary is + // aligned to `align` (both are powers of two and `align <= 4 KiB`). + if size < layout.align() { + size = layout.align(); + } + if size < MIN_BLOCK_SIZE { + size = MIN_BLOCK_SIZE; + } + if size > MAX_BLOCK_SIZE { + return None; + } + let shift = size.next_power_of_two().trailing_zeros(); + Some((shift - MIN_CLASS_SHIFT) as usize) +} diff --git a/crates/fspy_alloc/src/lib.rs b/crates/fspy_alloc/src/lib.rs new file mode 100644 index 000000000..90b2524cf --- /dev/null +++ b/crates/fspy_alloc/src/lib.rs @@ -0,0 +1,124 @@ +//! Lock-free, async-signal-safe global allocator for the fspy preload library. +//! +//! The preload library interposes libc functions that POSIX declares +//! async-signal-safe (`open`, `stat`, `execve`, ...). Programs may call these +//! from signal handlers, and — more commonly — from the child of `fork()` in a +//! multithreaded process, where only async-signal-safe calls are permitted: +//! the libc allocator's locks may be held forever by threads that no longer +//! exist after the fork. Routing the preload's Rust allocations through this +//! allocator keeps them safe in both contexts: +//! +//! - **No locks.** Every state transition is a lock-free compare-and-swap +//! loop: an attempt only retries because another running thread completed +//! its operation, so nothing ever waits on state that a thread which +//! vanished at `fork()` — or sits suspended under a signal handler — would +//! have to release. (Lock-free, not wait-free: an individual operation has +//! no fixed retry bound under active contention.) +//! - **No thread-locals.** TLS first-touch allocates through libc malloc on +//! some platforms (macOS thread-local variables), which would reintroduce +//! the hazard this crate exists to remove. +//! - **mmap-backed.** Memory comes straight from the kernel. On Linux the +//! allocator relies on nothing from libc: mapping syscalls are issued +//! directly (rustix's raw backend) and even the page size is discovered by +//! probing with raw syscalls. On macOS, which has no stable raw-syscall +//! ABI, calls go through the thin libSystem stubs. libc malloc is never +//! called anywhere. +//! +//! Design: power-of-two size classes (16 B ..= 64 KiB) carve blocks out of +//! 1 MiB slabs; freed blocks recycle through a per-class Treiber free list +//! made ABA-safe by a generation tag. Requests larger than the biggest class +//! (or over-aligned beyond 4 KiB) map and unmap directly. See the `pool` +//! module for the details. +//! +//! Because the allocator is a `const`-initialized static with no lazy setup, +//! it works from the very first allocation in the process — even before the +//! preload library's constructor runs. + +#![cfg_attr(not(test), no_std)] + +// Compile as an empty crate on non-unix targets: the allocator backs the unix +// preload library. A Windows backend can be added alongside `sys::Mmap` if +// the Windows preload ever needs one. + +#[cfg(unix)] +mod class; +#[cfg(unix)] +mod mapping; +#[cfg(unix)] +mod mmap; +#[cfg(unix)] +mod pool; +#[cfg(unix)] +mod slab; +#[cfg(unix)] +mod sys; + +#[cfg(unix)] +use core::{ + alloc::{GlobalAlloc, Layout}, + ptr::{self, NonNull}, +}; + +#[cfg(unix)] +use crate::{mmap::Mmap, pool::Pool}; + +/// A lock-free, async-signal-safe, fork-safe [`GlobalAlloc`] implementation. +/// +/// Intended to be installed as the `#[global_allocator]` of the fspy preload +/// library. All memory comes from anonymous mappings; libc malloc is never +/// called, no locks are taken, and no thread-local state is used. +/// +/// Capacity is bounded by design: each size class can hold at most 256 slabs +/// of 1 MiB (roughly 200 MiB per class). Requests beyond that — far outside +/// anything the preload library does — fail like any other out-of-memory +/// condition (`alloc` returns null). +#[cfg(unix)] +pub struct FspyAlloc { + pool: Pool, +} + +#[cfg(unix)] +impl FspyAlloc { + /// Creates the allocator. `const` so it can back a `static` with no + /// runtime initialization. + #[must_use] + pub const fn new() -> Self { + Self { pool: Pool::new() } + } +} + +#[cfg(unix)] +impl Default for FspyAlloc { + fn default() -> Self { + Self::new() + } +} + +// SAFETY: `Pool` hands out blocks that are non-null, at least `layout.size()` +// bytes large, aligned to at least `layout.align()`, and exclusively owned +// until returned via `dealloc`. Allocation failure is reported as null, and +// none of the methods unwind. +#[cfg(unix)] +unsafe impl GlobalAlloc for FspyAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + self.pool.alloc(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + } + + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + let Some(ptr) = NonNull::new(ptr) else { return }; + // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this + // allocator for this `layout`. + unsafe { self.pool.dealloc(ptr, layout) } + } + + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + self.pool.alloc_zeroed(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + } + + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + let Some(ptr) = NonNull::new(ptr) else { return ptr::null_mut() }; + // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this + // allocator for this `layout`, and `new_size` is non-zero. + unsafe { self.pool.realloc(ptr, layout, new_size) }.map_or(ptr::null_mut(), NonNull::as_ptr) + } +} diff --git a/crates/fspy_alloc/src/mapping.rs b/crates/fspy_alloc/src/mapping.rs new file mode 100644 index 000000000..d3f53e726 --- /dev/null +++ b/crates/fspy_alloc/src/mapping.rs @@ -0,0 +1,62 @@ +//! Owned memory regions obtained from a [`Sys`] provider. + +use core::{marker::PhantomData, mem, ptr::NonNull}; + +use crate::sys::Sys; + +/// An owned region obtained from `S`, released on drop. +/// +/// This is the only place that calls [`Sys::unmap`]: pool code either lets a +/// `Mapping` drop (probe scratch, install races, freed large allocations) or +/// deliberately leaks it with [`Mapping::into_raw`] (published slabs, live +/// large allocations). Reconstructing ownership from a raw pointer via +/// [`Mapping::from_raw`] is the single unsafe step. +pub struct Mapping { + ptr: NonNull, + size: usize, + align: usize, + sys: PhantomData S>, +} + +impl Mapping { + /// Maps `size` bytes of zero-initialized memory aligned to `align` + /// (a power of two). Returns `None` when memory is exhausted. + pub fn new(size: usize, align: usize) -> Option { + let ptr = S::map(size, align)?; + Some(Self { ptr, size, align, sys: PhantomData }) + } + + /// Reclaims ownership of a mapping previously released with + /// [`Mapping::into_raw`]. + /// + /// # Safety + /// + /// `ptr` must have come from `Mapping::::into_raw` (or `Sys::map`) + /// with exactly this `size` and `align`, the region must not be in use, + /// and ownership must not be reclaimed twice. + pub unsafe fn from_raw(ptr: NonNull, size: usize, align: usize) -> Self { + Self { ptr, size, align, sys: PhantomData } + } + + /// The mapped region's base address. + pub const fn ptr(&self) -> NonNull { + self.ptr + } + + /// Releases ownership without unmapping; the region lives until (unless) + /// [`Mapping::from_raw`] reclaims it. + pub const fn into_raw(self) -> NonNull { + let ptr = self.ptr; + mem::forget(self); + ptr + } +} + +impl Drop for Mapping { + fn drop(&mut self) { + // SAFETY: this type owns the mapping (constructed from `Sys::map` + // directly or via the `from_raw` contract), and after drop nothing + // can use it. + unsafe { S::unmap(self.ptr, self.size, self.align) } + } +} diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/fspy_alloc/src/mmap.rs new file mode 100644 index 000000000..053356b56 --- /dev/null +++ b/crates/fspy_alloc/src/mmap.rs @@ -0,0 +1,211 @@ +//! Kernel-backed [`Sys`] provider: anonymous mappings via direct syscalls. +//! +//! On Linux the provider relies on nothing from libc — rustix issues raw +//! syscalls, and even the page size is discovered with raw syscalls (see +//! [`page_size`]). macOS has no stable raw-syscall ABI, so calls go through +//! the thin libSystem stubs there. + +use core::{ + ptr::{self, NonNull}, + sync::atomic::{AtomicUsize, Ordering}, +}; + +use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; + +use crate::sys::Sys; + +/// Returns the kernel page size, or `None` when it cannot be determined. +/// +/// On Linux the value is discovered with raw syscalls only — no libc, no +/// `/proc`, no minimum kernel version (see [`query_page_size`]). On other +/// unix platforms (macOS, where every syscall goes through libSystem by +/// platform contract anyway) it comes from `sysconf` via rustix. Either +/// way the result is a process constant, validated as a power of two and +/// cached, so `map` and `unmap` can never disagree on rounding; an +/// undeterminable page size fails the allocation rather than guessing. +fn page_size() -> Option { + static CACHE: AtomicUsize = AtomicUsize::new(0); + let cached = CACHE.load(Ordering::Relaxed); + if cached != 0 { + return Some(cached); + } + let page = query_page_size()?; + if !page.is_power_of_two() { + return None; + } + // First store wins; every query returns the same value, so the cache + // only avoids repeated probing. + let _ = CACHE.compare_exchange(0, page, Ordering::Relaxed, Ordering::Relaxed); + Some(page) +} + +/// Discovers the page size by probing, using nothing but raw syscalls. +/// +/// `mprotect` fails with `EINVAL` unless its address is a multiple of the +/// page size, so re-protecting a scratch mapping at increasing +/// power-of-two offsets identifies the page size as the first offset the +/// kernel accepts (for a power-of-two page size P, the first power of two +/// P divides is P itself). A handful of syscalls, once per process: +/// async-signal-safe, fork-safe, and independent of libc, `/proc` +/// availability, and kernel version. +/// +/// Why not `rustix::param::page_size()` here (it's fine on macOS)? On +/// Linux the allocator must not rely on libc, which rules out rustix's +/// `use-libc-auxv` (`getauxval`) configuration — and rustix's libc-free +/// fallback is unusable *inside* a global allocator: +/// +/// - Its lazy init tries `prctl(PR_GET_AUXV)` (kernel 6.4+ only) and +/// otherwise reads `/proc/self/auxv`; with rustix's `alloc` feature +/// enabled, that read path heap-allocates (`Vec`) — through *this* +/// allocator, whose `map` is the caller waiting on the page size — +/// recursing unboundedly. And we cannot pin `alloc` off: Cargo feature +/// unification lets any other rustix user in the build graph enable it +/// for our copy. +/// - It panics on read errors, truncated auxv, or both sources being +/// unavailable, where this allocator requires failure to surface as +/// `None` (panic formatting itself allocates, re-entering the same +/// uninitialized path). +/// +/// The probe has neither problem: no allocation, no panic, no minimum +/// kernel, and its worst case is `None`. +#[cfg(target_os = "linux")] +fn query_page_size() -> Option { + use rustix::mm::{MprotectFlags, mprotect}; + + /// Smallest page size of any Linux configuration; the probe starts + /// here. + const MIN_PROBE_PAGE: usize = 4096; + /// Generous upper bound for the probe; no supported configuration + /// uses larger pages. + const MAX_PROBE_PAGE: usize = 1 << 20; + + // Large enough that every probe below stays inside the mapping even + // after the kernel rounds the one-byte length up to a full page. + let scratch = map_anonymous(2 * MAX_PROBE_PAGE)?; + let mut page = None; + let mut offset = MIN_PROBE_PAGE; + while offset <= MAX_PROBE_PAGE { + // SAFETY: `scratch + offset` (plus the page-rounded single byte) + // lies within the scratch mapping, which we exclusively own; the + // protection flags match the mapping's existing ones. + let accepted = unsafe { + mprotect( + scratch.as_ptr().add(offset).cast(), + 1, + MprotectFlags::READ | MprotectFlags::WRITE, + ) + } + .is_ok(); + if accepted { + page = Some(offset); + break; + } + offset *= 2; + } + // SAFETY: releasing the scratch mapping created above. + let _ = unsafe { munmap(scratch.as_ptr().cast(), 2 * MAX_PROBE_PAGE) }; + page +} + +#[cfg(not(target_os = "linux"))] +#[expect( + clippy::unnecessary_wraps, + reason = "must match the signature of the fallible Linux probe variant" +)] +fn query_page_size() -> Option { + Some(rustix::param::page_size()) +} + +const fn round_up(value: usize, align: usize) -> usize { + (value + align - 1) & !(align - 1) +} + +fn map_anonymous(len: usize) -> Option> { + // SAFETY: a fresh anonymous private mapping at no particular address + // has no memory-safety preconditions. + let ptr = unsafe { + mmap_anonymous(ptr::null_mut(), len, ProtFlags::READ | ProtFlags::WRITE, MapFlags::PRIVATE) + } + .ok()?; + NonNull::new(ptr.cast::()) +} + +/// Kernel-backed provider: anonymous mappings obtained through direct +/// syscalls, alignment achieved by over-mapping and trimming. +pub struct Mmap; + +impl Sys for Mmap { + fn map(size: usize, align: usize) -> Option> { + debug_assert!(align.is_power_of_two()); + let page = page_size()?; + let size = round_up(size.max(1), page); + + if align <= page { + // Mapping results are aligned to the (verified real) page + // size. + return map_anonymous(size); + } + + // Over-map by `align`, then unmap the misaligned head and the + // leftover tail. All cut points are page-aligned: `raw` and + // `aligned` are page-aligned, and `size`/`align` are multiples of + // the page size. + let raw = map_anonymous(size.checked_add(align)?)?; + let raw_addr = raw.as_ptr().addr(); + let aligned_addr = round_up(raw_addr, align); + let head = aligned_addr - raw_addr; + let tail = align - head; + if head > 0 { + // SAFETY: `[raw_addr, raw_addr + head)` lies within the fresh + // mapping and `raw_addr` is page-aligned. Failure is + // impossible for a region we own; if it happened anyway the + // pages would merely stay mapped. + let _ = unsafe { munmap(raw.as_ptr().cast(), head) }; + } + if tail > 0 { + let tail_start = raw.as_ptr().with_addr(aligned_addr + size); + // SAFETY: `[aligned_addr + size, raw_addr + size + align)` + // lies within the fresh mapping and its start is page-aligned + // (see above). Failure is impossible for a region we own. + let _ = unsafe { munmap(tail_start.cast(), tail) }; + } + // `aligned_addr` lies inside a successful mapping and so can + // never be zero, but checked construction costs nothing here. + core::num::NonZero::new(aligned_addr).map(|addr| raw.with_addr(addr)) + } + + unsafe fn unmap(ptr: NonNull, size: usize, _align: usize) { + // A successful `map` proved the page size, so this cannot fail + // for a live mapping; if it somehow did, leaking the region is + // the only safe response. + let Some(page) = page_size() else { return }; + // Whether or not `map` trimmed for alignment, the retained region + // is exactly `[ptr, ptr + round_up(size, page))`. + let size = round_up(size.max(1), page); + // SAFETY: caller contract — `ptr`/`size` describe a live mapping + // returned by `map`. Failure is impossible for a region we own. + let _ = unsafe { munmap(ptr.as_ptr().cast(), size) }; + } +} + +#[cfg(all(test, not(miri)))] +mod tests { + use super::Mmap; + use crate::sys::Sys; + + #[test] + fn real_mappings_are_aligned_zeroed_and_writable() { + for (size, align) in [(1, 1), (4096, 4096), (100, 1 << 20), (5 << 20, 4096)] { + let ptr = Mmap::map(size, align).unwrap(); + assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); + for i in 0..size { + // SAFETY: fresh exclusive mapping of at least `size` bytes. + assert_eq!(unsafe { ptr.as_ptr().add(i).read() }, 0); + } + // SAFETY: fresh exclusive mapping of at least `size` bytes. + unsafe { ptr.as_ptr().write_bytes(0x5A, size) }; + // SAFETY: mapped above with the same size and alignment. + unsafe { Mmap::unmap(ptr, size, align) }; + } + } +} diff --git a/crates/fspy_alloc/src/pool.rs b/crates/fspy_alloc/src/pool.rs new file mode 100644 index 000000000..6bb4f82f1 --- /dev/null +++ b/crates/fspy_alloc/src/pool.rs @@ -0,0 +1,705 @@ +//! Lock-free size-class pool. +//! +//! Layouts round up to a power-of-two class (see the `class` module) whose +//! blocks are carved from slabs (see the `slab` module for the memory +//! layout). +//! +//! # Concurrency +//! +//! Each class has a Treiber-stack free list plus a bump cursor over the +//! newest ("active") slab: +//! +//! - **alloc** pops the free list, or carves the next block off the active +//! slab, installing a fresh slab when the active one is exhausted. +//! - **dealloc** pushes the block back onto its class's free list. +//! +//! The 64-bit list head packs a 40-bit generation tag next to the 24-bit +//! block reference; the tag advances on every successful push and pop, which +//! makes the classic Treiber-stack ABA failure require a thread to stall +//! between its head load and CAS while other threads perform an exact +//! multiple of 2^40 head mutations. That is a probabilistic defense, not a +//! formal impossibility — see [`HEAD_TAG_SHIFT`]. Slab installation races +//! are resolved with a compare-and-swap on the slab table slot; the loser +//! simply returns its (never-published) mapping. +//! +//! The pool is **lock-free, not wait-free**: a CAS loop can retry +//! indefinitely under contention, but a retry only ever happens because +//! another thread completed an operation, so system-wide progress never +//! stalls — and, the property fork- and signal-safety actually require, no +//! operation ever waits on state that only a suspended or vanished thread +//! could release. Blocks are aligned to `min(block size, 4 KiB)`; requests +//! with stricter alignment (or size beyond the largest class) bypass the +//! pool and map directly. + +use core::{ + alloc::Layout, + marker::PhantomData, + ptr, + ptr::NonNull, + sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, Ordering}, +}; + +use crate::{ + class::{CLASS_COUNT, class_of}, + mapping::Mapping, + slab::{BLOCKS_PER_SLAB, SLAB_SIZE, Slab}, + sys::Sys, +}; + +/// Bounded by the 8 bits reserved for slab indices in a packed block +/// reference. Caps each class at ~200 MiB. +const MAX_SLABS_PER_CLASS: usize = 256; + +/// Sentinel for "no block" in the 24-bit packed-reference field of a +/// free-list head. Never collides with a real reference: block indices stay +/// below `0xFFFF` (asserted with the slab geometry below). +const NO_BLOCK: u32 = 0x00FF_FFFF; +/// Sentinel for "no slab installed yet" in [`ClassState::active`]. +const NO_SLAB: u32 = u32::MAX; + +const _: () = { + assert!(MAX_SLABS_PER_CLASS <= 1 << 8, "slab index must fit in 8 bits"); + let mut class = 0; + while class < CLASS_COUNT { + assert!( + BLOCKS_PER_SLAB[class] <= 0xFFFF, + "block indices must fit in 16 bits and stay below the NO_BLOCK sentinel" + ); + class += 1; + } +}; + +/// Packs a block's location into 24 bits: slab index in bits 16..24, block +/// index in bits 0..16. +#[expect( + clippy::cast_possible_truncation, + reason = "callers pass indices bounded by MAX_SLABS_PER_CLASS and BLOCKS_PER_SLAB" +)] +const fn pack_ref(slab_idx: usize, block_idx: usize) -> u32 { + ((slab_idx as u32) << 16) | (block_idx as u32) +} + +const fn unpack_ref(packed: u32) -> (usize, usize) { + (((packed >> 16) & 0xFF) as usize, (packed & 0xFFFF) as usize) +} + +/// A free-list head is `[generation tag : 40 | packed block reference : 24]`. +/// The tag advances on every successful push and pop, so a stale +/// compare-and-swap can only succeed if its thread stalls between head load +/// and CAS while others perform an exact multiple of 2^40 head mutations — +/// not a formal impossibility, but hours of maximum-rate churn inside one +/// stalled instruction window. +const HEAD_TAG_SHIFT: u32 = 24; +const HEAD_TAG_MASK: u64 = (1 << 40) - 1; +const HEAD_REF_MASK: u64 = (1 << HEAD_TAG_SHIFT) - 1; + +/// Splits a free-list head into `(generation tag, packed block reference)`. +const fn head_parts(head: u64) -> (u64, u32) { + // The mask keeps the reference within 24 bits, so the cast is lossless. + (head >> HEAD_TAG_SHIFT, (head & HEAD_REF_MASK) as u32) +} + +const fn head_from_parts(tag: u64, block_ref: u32) -> u64 { + ((tag & HEAD_TAG_MASK) << HEAD_TAG_SHIFT) | block_ref as u64 +} + +/// Aligned to its own cache-line region so hot heads of different classes +/// don't false-share. +#[repr(align(128))] +struct ClassState { + /// Treiber free-list head: `[generation tag : 40 | packed ref : 24]`. + head: AtomicU64, + /// Index of the slab currently being carved, or [`NO_SLAB`] before the + /// first slab is installed. + active: AtomicU32, + /// Base addresses of installed slabs. Written once (null → mapping) and + /// never cleared. + slabs: [AtomicPtr; MAX_SLABS_PER_CLASS], +} + +impl ClassState { + const fn new() -> Self { + Self { + head: AtomicU64::new(head_from_parts(0, NO_BLOCK)), + active: AtomicU32::new(NO_SLAB), + slabs: [const { AtomicPtr::new(ptr::null_mut()) }; MAX_SLABS_PER_CLASS], + } + } +} + +/// A block handed out by [`Pool::alloc_block`], remembering whether it came +/// off the free list (may contain stale data) or was carved fresh off a slab +/// (still kernel-zeroed). +enum ClassBlock { + Recycled(NonNull), + Carved(NonNull), +} + +enum Carve { + Block(NonNull), + /// A new slab was (or concurrently got) installed; retry the allocation. + Retry, + /// Out of slab table entries or out of memory. + Exhausted, +} + +/// The allocator core, generic over its [`Sys`] memory provider. +pub struct Pool { + classes: [ClassState; CLASS_COUNT], + sys: PhantomData S>, +} + +impl Pool { + #[expect( + clippy::large_stack_arrays, + reason = "the class table (~30 KiB) is only ever materialized into a const-initialized static, never built on a runtime stack" + )] + pub const fn new() -> Self { + Self { classes: [const { ClassState::new() }; CLASS_COUNT], sys: PhantomData } + } + + /// Allocates memory for `layout`. Returns `None` on exhaustion. + pub fn alloc(&self, layout: Layout) -> Option> { + match class_of(layout) { + Some(class) => match self.alloc_block(class)? { + ClassBlock::Recycled(ptr) | ClassBlock::Carved(ptr) => Some(ptr), + }, + // Deliberately leaked until `dealloc` reclaims ownership. + None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), + } + } + + /// Like [`Pool::alloc`], but the returned memory is zeroed. + pub fn alloc_zeroed(&self, layout: Layout) -> Option> { + match class_of(layout) { + Some(class) => match self.alloc_block(class)? { + // Freshly carved blocks are still kernel-zeroed. + ClassBlock::Carved(ptr) => Some(ptr), + ClassBlock::Recycled(ptr) => { + // SAFETY: the block is freshly allocated, exclusively + // ours, and at least `layout.size()` bytes. + unsafe { ptr.as_ptr().write_bytes(0, layout.size()) }; + Some(ptr) + } + }, + // Fresh mappings are zeroed by the kernel; deliberately leaked + // until `dealloc` reclaims ownership. + None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), + } + } + + /// Releases memory obtained from this pool. + /// + /// # Safety + /// + /// `ptr` must have been returned by this pool for exactly this `layout` + /// and must not be used afterwards. + pub unsafe fn dealloc(&self, ptr: NonNull, layout: Layout) { + match class_of(layout) { + // SAFETY: caller contract — a `Some(class)` layout was served + // from the pool, so `ptr` is a live block of this class. + Some(class) => unsafe { self.push_free(ptr, class) }, + // SAFETY: caller contract — a `None` layout was served by the + // mapping path with these parameters and is no longer in use, so + // ownership can be reclaimed (and the region dropped). + None => drop(unsafe { Mapping::::from_raw(ptr, layout.size(), layout.align()) }), + } + } + + /// Grows or shrinks an allocation, preserving contents up to the smaller + /// of the old and new sizes. Returns `None` on exhaustion (the original + /// allocation stays valid). + /// + /// # Safety + /// + /// `ptr` must have been returned by this pool for exactly `layout`, and + /// `new_size` must be non-zero. + pub unsafe fn realloc( + &self, + ptr: NonNull, + layout: Layout, + new_size: usize, + ) -> Option> { + let new_layout = Layout::from_size_align(new_size, layout.align()).ok()?; + let class = class_of(layout); + if class.is_some() && class == class_of(new_layout) { + // Same size class: the existing block already fits. + return Some(ptr); + } + let new_ptr = self.alloc(new_layout)?; + // SAFETY: `new_ptr` is a fresh exclusive allocation of at least + // `new_size` bytes; `ptr` is valid for `layout.size()` bytes (caller + // contract); distinct allocations never overlap. + unsafe { + new_ptr.as_ptr().copy_from_nonoverlapping(ptr.as_ptr(), layout.size().min(new_size)); + } + // SAFETY: caller contract — `ptr` came from this pool with `layout`. + unsafe { self.dealloc(ptr, layout) }; + Some(new_ptr) + } + + fn alloc_block(&self, class: usize) -> Option { + loop { + if let Some(ptr) = self.pop_free(class) { + return Some(ClassBlock::Recycled(ptr)); + } + match self.carve(class) { + Carve::Block(ptr) => return Some(ClassBlock::Carved(ptr)), + Carve::Retry => {} + Carve::Exhausted => return None, + } + } + } + + /// Returns the published slab of `class` at `slab_idx`, if installed. + fn published_slab(&self, class: usize, slab_idx: usize) -> Option { + let base = NonNull::new(self.classes[class].slabs[slab_idx].load(Ordering::Acquire))?; + // SAFETY: non-null slab-table entries are only published (with + // `Release`, paired with the `Acquire` above) after + // `Slab::init_header` ran on a fresh SLAB_SIZE-aligned mapping, and + // are never cleared or unmapped while the pool is in use. + Some(unsafe { Slab::from_published(base, class) }) + } + + /// Pops a block off the class's free list. + fn pop_free(&self, class: usize) -> Option> { + let state = &self.classes[class]; + loop { + let head = state.head.load(Ordering::Acquire); + let (tag, block_ref) = head_parts(head); + if block_ref == NO_BLOCK { + return None; + } + let (slab_idx, block_idx) = unpack_ref(block_ref); + // A listed block was carved from its slab, so the slab is always + // published; `None` here is unreachable in practice. + let slab = self.published_slab(class, slab_idx)?; + // Read the successor link before the CAS. If another thread pops + // this block first, the tag comparison below fails and the value + // read here is discarded; since links live in the atomic side + // table, the racing read itself is well-defined. + let next = slab.link(block_idx).load(Ordering::Relaxed); + let new_head = head_from_parts(tag.wrapping_add(1), next); + if state + .head + .compare_exchange_weak(head, new_head, Ordering::AcqRel, Ordering::Relaxed) + .is_ok() + { + return Some(slab.block(block_idx)); + } + } + } + + /// Pushes a block onto the class's free list. + /// + /// # Safety + /// + /// `ptr` must be a block of `class` previously returned by this pool and + /// no longer in use. + unsafe fn push_free(&self, ptr: NonNull, class: usize) { + // SAFETY: caller contract — `ptr` is a live block of `class` from + // this pool. (The impossible `None` would merely leak the block.) + let Some(slab) = (unsafe { Slab::of_block(ptr, class) }) else { return }; + let block_idx = slab.block_index(ptr); + let packed = pack_ref(slab.slab_idx(), block_idx); + let link = slab.link(block_idx); + let state = &self.classes[class]; + loop { + let head = state.head.load(Ordering::Relaxed); + let (tag, head_ref) = head_parts(head); + link.store(head_ref, Ordering::Relaxed); + let new_head = head_from_parts(tag.wrapping_add(1), packed); + // `Release` publishes the link store above to the eventual popper. + if state + .head + .compare_exchange_weak(head, new_head, Ordering::Release, Ordering::Relaxed) + .is_ok() + { + return; + } + } + } + + /// Carves the next block off the class's active slab, installing a new + /// slab if the active one is exhausted (or none exists yet). + fn carve(&self, class: usize) -> Carve { + let state = &self.classes[class]; + let active = state.active.load(Ordering::Acquire); + // `active` is only published after its slab pointer, so a live + // `active` always resolves to a published slab. + if active != NO_SLAB + && let Some(slab) = self.published_slab(class, active as usize) + { + let carved_idx = slab.carved().fetch_add(1, Ordering::Relaxed); + if let Some(block_idx) = carved_to_block_idx(carved_idx, class) { + return Carve::Block(slab.block(block_idx)); + } + // Active slab exhausted; fall through to install the next one. + } + let next_idx = if active == NO_SLAB { 0 } else { active as usize + 1 }; + if next_idx >= MAX_SLABS_PER_CLASS { + return Carve::Exhausted; + } + self.install_slab(class, next_idx); + if self.published_slab(class, next_idx).is_none() { + // Our mapping failed and no other thread succeeded either. + return Carve::Exhausted; + } + #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] + let next_active = next_idx as u32; + // Advance `active`; losing the race just means another thread already + // advanced it. Either way the retry re-reads it. + let _ = + state.active.compare_exchange(active, next_active, Ordering::AcqRel, Ordering::Relaxed); + Carve::Retry + } + + /// Maps and publishes the slab at `slab_idx`, unless another thread beats + /// us to it (or the mapping fails, leaving the slot null). + fn install_slab(&self, class: usize, slab_idx: usize) { + let state = &self.classes[class]; + if !state.slabs[slab_idx].load(Ordering::Acquire).is_null() { + return; + } + let Some(mapping) = Mapping::::new(SLAB_SIZE, SLAB_SIZE) else { return }; + #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] + let idx = slab_idx as u32; + // SAFETY: `mapping` is a fresh, exclusive, zero-initialized, + // SLAB_SIZE-byte and SLAB_SIZE-aligned mapping. + unsafe { Slab::init_header(mapping.ptr(), idx) }; + if state.slabs[slab_idx] + .compare_exchange( + ptr::null_mut(), + mapping.ptr().as_ptr(), + Ordering::Release, + Ordering::Relaxed, + ) + .is_ok() + { + // Published: the slab now lives for the rest of the process. + let _ = mapping.into_raw(); + } + // Otherwise another thread installed this slot first; our + // never-published mapping is released when `mapping` drops. + } + + /// Tears down all slab mappings. Test-only: the global allocator lives in + /// a static and never releases its slabs, but tests (and Miri's leak + /// checker) want a clean shutdown. + #[cfg(test)] + fn unmap_all_slabs(&mut self) { + for state in &mut self.classes { + *state.head.get_mut() = head_from_parts(0, NO_BLOCK); + *state.active.get_mut() = NO_SLAB; + for slot in &mut state.slabs { + let slab = core::mem::replace(slot.get_mut(), ptr::null_mut()); + if let Some(base) = NonNull::new(slab) { + // SAFETY: `base` was leaked into the table by + // `install_slab` with these parameters; `&mut self` + // guarantees no concurrent (or future) use of its blocks. + drop(unsafe { Mapping::::from_raw(base, SLAB_SIZE, SLAB_SIZE) }); + } + } + } + } +} + +/// Converts a raw carve-counter value into a block index, or `None` if the +/// slab is exhausted. +fn carved_to_block_idx(carved: u64, class: usize) -> Option { + let idx = usize::try_from(carved).ok()?; + (idx < BLOCKS_PER_SLAB[class]).then_some(idx) +} + +#[cfg(test)] +mod tests { + use std::{boxed::Box, sync::mpsc, thread, vec::Vec}; + + use super::*; + use crate::class::{MAX_BLOCK_SIZE, MAX_POOL_ALIGN, block_size}; + + /// Host-allocator-backed [`Sys`] so the pool core runs under Miri and on + /// any platform. + struct TestSys; + + impl Sys for TestSys { + fn map(size: usize, align: usize) -> Option> { + let layout = Layout::from_size_align(size.max(1), align).ok()?; + // SAFETY: `layout` has non-zero size. + NonNull::new(unsafe { std::alloc::alloc_zeroed(layout) }) + } + + unsafe fn unmap(ptr: NonNull, size: usize, align: usize) { + let layout = Layout::from_size_align(size.max(1), align).unwrap(); + // SAFETY: caller contract — `ptr` was returned by `map`, which + // used exactly this layout. + unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) } + } + } + + fn with_pool(test: impl FnOnce(&Pool)) { + let mut pool = Box::new(Pool::::new()); + test(&pool); + pool.unmap_all_slabs(); + } + + fn layout(size: usize, align: usize) -> Layout { + Layout::from_size_align(size, align).unwrap() + } + + #[test] + fn head_tag_wraps_within_its_field() { + let head = head_from_parts(HEAD_TAG_MASK, 42); + assert_eq!(head_parts(head), (HEAD_TAG_MASK, 42)); + // Advancing the maximal tag must wrap to zero without touching the + // reference bits. + let wrapped = head_from_parts(HEAD_TAG_MASK.wrapping_add(1), 42); + assert_eq!(head_parts(wrapped), (0, 42)); + } + + #[test] + fn packed_refs_cannot_collide_with_the_sentinel() { + for &blocks in &BLOCKS_PER_SLAB { + let max_ref = pack_ref(MAX_SLABS_PER_CLASS - 1, blocks - 1); + assert_ne!(max_ref, NO_BLOCK); + assert!(u64::from(max_ref) <= HEAD_REF_MASK); + } + } + + #[test] + fn classify_boundaries() { + assert_eq!(class_of(layout(1, 1)), Some(0)); + assert_eq!(class_of(layout(16, 1)), Some(0)); + assert_eq!(class_of(layout(17, 1)), Some(1)); + assert_eq!(class_of(layout(MAX_BLOCK_SIZE, 8)), Some(CLASS_COUNT - 1)); + assert_eq!(class_of(layout(MAX_BLOCK_SIZE + 1, 8)), None); + // Alignment can raise the class. + assert_eq!(class_of(layout(8, 1024)), class_of(layout(1024, 8))); + // Over-aligned layouts bypass the pool. + assert_eq!(class_of(layout(16, MAX_POOL_ALIGN * 2)), None); + } + + #[test] + fn round_trip_every_class() { + with_pool(|pool| { + for class in 0..CLASS_COUNT { + let l = layout(block_size(class), 8); + let ptr = pool.alloc(l).unwrap(); + // SAFETY: fresh exclusive allocation of `block_size` bytes. + unsafe { ptr.as_ptr().write_bytes(0xAB, l.size()) }; + // SAFETY: reading back the block we just wrote. + let last = unsafe { ptr.as_ptr().add(l.size() - 1).read() }; + assert_eq!(last, 0xAB); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn block_alignment() { + with_pool(|pool| { + for align in [8, 64, 1024, MAX_POOL_ALIGN] { + let l = layout(24, align); + let ptr = pool.alloc(l).unwrap(); + assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn free_list_recycles_lifo() { + with_pool(|pool| { + let l = layout(100, 8); + let first = pool.alloc(l).unwrap(); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(first, l) }; + let second = pool.alloc(l).unwrap(); + assert_eq!(first, second); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(second, l) }; + }); + } + + #[test] + fn carves_across_multiple_slabs() { + with_pool(|pool| { + // The largest class has the fewest blocks per slab, so a couple + // dozen live allocations force several slab installations. + let l = layout(MAX_BLOCK_SIZE, 8); + let count = BLOCKS_PER_SLAB[CLASS_COUNT - 1] * 3 + 1; + let blocks: Vec> = (0..count).map(|_| pool.alloc(l).unwrap()).collect(); + for (i, ptr) in blocks.iter().enumerate() { + assert!(blocks[..i].iter().all(|other| other != ptr), "duplicate block"); + // SAFETY: live exclusive allocation of `MAX_BLOCK_SIZE` bytes. + unsafe { ptr.as_ptr().write_bytes(0x5A, l.size()) }; + } + for ptr in blocks { + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn large_allocations_bypass_pool() { + with_pool(|pool| { + for l in [layout(MAX_BLOCK_SIZE + 1, 8), layout(5 << 20, 8), layout(64, 8192)] { + let ptr = pool.alloc(l).unwrap(); + assert_eq!(ptr.as_ptr().addr() % l.align(), 0); + // SAFETY: fresh exclusive allocation of `l.size()` bytes. + unsafe { ptr.as_ptr().write_bytes(0xCD, l.size()) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + + #[test] + fn realloc_within_class_keeps_block() { + with_pool(|pool| { + let l = layout(100, 8); + let ptr = pool.alloc(l).unwrap(); + // SAFETY: realloc contract — `ptr` allocated with `l`. + let grown = unsafe { pool.realloc(ptr, l, 120) }.unwrap(); + assert_eq!(ptr, grown, "same class must realloc in place"); + // SAFETY: allocated above; 120 rounds to the same class as 100. + unsafe { pool.dealloc(grown, layout(120, 8)) }; + }); + } + + #[test] + fn realloc_across_classes_preserves_contents() { + with_pool(|pool| { + let l = layout(64, 8); + let ptr = pool.alloc(l).unwrap(); + for i in 0..64u8 { + // SAFETY: live exclusive allocation of 64 bytes. + unsafe { ptr.as_ptr().add(usize::from(i)).write(i) }; + } + // SAFETY: realloc contract — `ptr` allocated with `l`. + let grown = unsafe { pool.realloc(ptr, l, 4096) }.unwrap(); + for i in 0..64u8 { + // SAFETY: `grown` is live for 4096 bytes. + let got = unsafe { grown.as_ptr().add(usize::from(i)).read() }; + assert_eq!(got, i); + } + // Shrink across classes, and from the large path back into a class. + // SAFETY: `grown` allocated with the 4096 layout above. + let shrunk = unsafe { pool.realloc(grown, layout(4096, 8), 8) }.unwrap(); + // SAFETY: `shrunk` is live for 8 bytes. + assert_eq!(unsafe { shrunk.as_ptr().read() }, 0); + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(shrunk, layout(8, 8)) }; + }); + } + + #[test] + fn alloc_zeroed_scrubs_recycled_blocks() { + with_pool(|pool| { + let l = layout(256, 8); + let dirty = pool.alloc(l).unwrap(); + // SAFETY: live exclusive allocation of 256 bytes. + unsafe { dirty.as_ptr().write_bytes(0xFF, l.size()) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(dirty, l) }; + + let zeroed = pool.alloc_zeroed(l).unwrap(); + assert_eq!( + zeroed, dirty, + "must recycle the dirty block for this test to be meaningful" + ); + for i in 0..l.size() { + // SAFETY: live exclusive allocation of 256 bytes. + assert_eq!(unsafe { zeroed.as_ptr().add(i).read() }, 0); + } + // SAFETY: allocated above with the same layout. + unsafe { pool.dealloc(zeroed, l) }; + }); + } + + /// Cheap deterministic PRNG so the stress test needs no dependencies. + fn lcg(state: &mut u64) -> u64 { + *state = + state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407); + *state >> 33 + } + + /// An owned allocation in flight between threads. + struct SendBlock(NonNull, Layout); + // SAFETY: SendBlock represents exclusive ownership of the block, which is + // transferred wholesale to the receiving thread. + unsafe impl Send for SendBlock {} + + #[test] + fn concurrent_stress_with_cross_thread_frees() { + const THREADS: usize = if cfg!(miri) { 3 } else { 8 }; + const OPS: usize = if cfg!(miri) { 60 } else { 20_000 }; + + with_pool(|pool| { + thread::scope(|scope| { + // Each thread frees blocks allocated by its neighbor, so + // pushes and pops of the same list run on different threads. + // Bounded channels keep the number of in-flight blocks well + // under the pool's per-class capacity regardless of thread + // scheduling. + let (senders, receivers): (Vec<_>, Vec<_>) = + (0..THREADS).map(|_| mpsc::sync_channel::(64)).unzip(); + let mut senders_rotated: Vec<_> = senders.into_iter().map(Some).collect(); + senders_rotated.rotate_left(1); + + for (thread_idx, (receiver, sender)) in + receivers.into_iter().zip(&mut senders_rotated).enumerate() + { + let sender = sender.take().unwrap(); + let pool = &*pool; + scope.spawn(move || { + let mut rng = 0x9E37_79B9_7F4A_7C15_u64 ^ thread_idx as u64; + for _ in 0..OPS { + let size = 1 + usize::try_from(lcg(&mut rng)).unwrap() + % (3 * MAX_BLOCK_SIZE / 2); + let align = 1 << (lcg(&mut rng) % 7); + let l = layout(size, align); + let ptr = pool.alloc(l).unwrap(); + let marker = u8::try_from(lcg(&mut rng) & 0xFF).unwrap(); + // SAFETY: fresh exclusive allocation of `size` bytes. + unsafe { + ptr.as_ptr().write(marker); + ptr.as_ptr().add(size - 1).write(marker); + } + // SAFETY: we still exclusively own the block. + let (first, last) = + unsafe { (ptr.as_ptr().read(), ptr.as_ptr().add(size - 1).read()) }; + assert_eq!((first, last), (marker, marker), "block corrupted"); + if let Err(returned) = sender.try_send(SendBlock(ptr, l)) { + // Neighbor's queue is full (or it finished); + // free locally instead of blocking, which in + // a ring of senders could deadlock. + let (mpsc::TrySendError::Full(SendBlock(ptr, l)) + | mpsc::TrySendError::Disconnected(SendBlock(ptr, l))) = returned; + // SAFETY: allocated above with layout `l`. + unsafe { pool.dealloc(ptr, l) }; + } + // Drain what our own producer has sent so far, + // interleaving cross-thread pops with pushes. + while let Ok(SendBlock(ptr, l)) = receiver.try_recv() { + // SAFETY: the neighbor allocated `ptr` with + // `l` and transferred ownership over the + // channel. + unsafe { pool.dealloc(ptr, l) }; + } + } + drop(sender); + for SendBlock(ptr, l) in receiver { + // SAFETY: the neighbor allocated `ptr` with `l` + // and transferred ownership over the channel. + unsafe { pool.dealloc(ptr, l) }; + } + }); + } + }); + }); + } +} diff --git a/crates/fspy_alloc/src/slab.rs b/crates/fspy_alloc/src/slab.rs new file mode 100644 index 000000000..3aca476dc --- /dev/null +++ b/crates/fspy_alloc/src/slab.rs @@ -0,0 +1,203 @@ +//! Slab memory layout and the [`Slab`] view type. +//! +//! Each size class carves fixed-size blocks out of [`SLAB_SIZE`]-byte, +//! `SLAB_SIZE`-aligned slabs. A slab looks like this: +//! +//! ```text +//! | SlabHeader | links: [AtomicU32; blocks] | pad to 4 KiB | block 0 | block 1 | ... | +//! ``` +//! +//! Free blocks are chained through the `links` side table — never through +//! block memory itself — so block payloads are only ever plain data. That +//! keeps every concurrent access well-defined under the memory model (no +//! mixed atomic/non-atomic reads of the same bytes) and Miri-clean. + +use core::{ + num::NonZero, + ptr::NonNull, + sync::atomic::{AtomicU32, AtomicU64}, +}; + +use crate::class::{CLASS_COUNT, MAX_POOL_ALIGN, block_size}; + +/// Slab size and alignment. Masking a block address with `!(SLAB_SIZE - 1)` +/// recovers its slab base. +pub const SLAB_SIZE: usize = 1 << 20; + +const fn round_up(value: usize, align: usize) -> usize { + (value + align - 1) & !(align - 1) +} + +/// Byte offset of the block area inside a slab holding `blocks` blocks. +const fn block_area_offset(blocks: usize) -> usize { + round_up(size_of::() + blocks * size_of::(), MAX_POOL_ALIGN) +} + +const fn compute_blocks_per_slab(class: usize) -> usize { + let size = block_size(class); + // Upper bound ignoring metadata, then shrink until header, link table, + // padding, and blocks all fit in one slab. + let mut blocks = SLAB_SIZE / size; + while block_area_offset(blocks) + blocks * size > SLAB_SIZE { + blocks -= 1; + } + blocks +} + +/// Blocks per slab, for each size class. +pub const BLOCKS_PER_SLAB: [usize; CLASS_COUNT] = { + let mut table = [0usize; CLASS_COUNT]; + let mut class = 0; + while class < CLASS_COUNT { + table[class] = compute_blocks_per_slab(class); + class += 1; + } + table +}; + +const _: () = { + assert!(size_of::() == 16); + assert!(MAX_POOL_ALIGN >= align_of::()); + let mut class = 0; + while class < CLASS_COUNT { + assert!(BLOCKS_PER_SLAB[class] > 0); + assert!( + block_area_offset(BLOCKS_PER_SLAB[class]) + BLOCKS_PER_SLAB[class] * block_size(class) + <= SLAB_SIZE + ); + class += 1; + } +}; + +/// Lives at the base of every slab mapping, ahead of the link table. +#[repr(C)] +struct SlabHeader { + /// This slab's index in its class's slab table. + slab_idx: u32, + _reserved: u32, + /// Number of blocks ever carved off this slab (monotonic; values at or + /// beyond the class's blocks-per-slab mean the slab is exhausted). 64-bit + /// so over-counting by racing threads can never wrap it around. + carved: AtomicU64, +} + +/// A view of one live slab mapping of a particular class. +/// +/// # Invariant +/// +/// `base` points to a [`SLAB_SIZE`]-byte, `SLAB_SIZE`-aligned mapping laid +/// out for `class` (initialized header, link table, block area) that is +/// never unmapped while the pool is in use. All raw-pointer arithmetic on +/// slab memory lives in this type's methods; the unsafe constructors are the +/// only places the invariant is asserted, and every accessor then relies on +/// it for bounds and liveness. +#[derive(Clone, Copy)] +pub struct Slab { + base: NonNull, + class: usize, +} + +impl Slab { + /// Wraps a slab pointer loaded from a class's slab table. + /// + /// # Safety + /// + /// `base` must be a non-null entry of the class's slab table. Such + /// entries are only published after [`Slab::init_header`] ran on a + /// suitable mapping, and are never cleared or unmapped while the pool is + /// in use, so the type invariant holds. + pub const unsafe fn from_published(base: NonNull, class: usize) -> Self { + Self { base, class } + } + + /// Recovers the slab containing a live pool block: slabs are + /// `SLAB_SIZE`-aligned, so masking the block address's low bits yields + /// the slab base (`with_addr` keeps the block pointer's provenance, which + /// covers the whole slab mapping it was carved from). Returns `None` only + /// for an address whose masked base would be null, which no real block + /// can produce. + /// + /// # Safety + /// + /// `block` must be a block of `class` previously handed out by this pool + /// and thus carved from a published slab of this class. + pub unsafe fn of_block(block: NonNull, class: usize) -> Option { + let base_addr = NonZero::new(block.as_ptr().addr() & !(SLAB_SIZE - 1))?; + Some(Self { base: block.with_addr(base_addr), class }) + } + + /// Initializes the header of a fresh slab mapping, making it publishable + /// into a slab table. + /// + /// # Safety + /// + /// `base` must point to a fresh, exclusive, zero-initialized, + /// `SLAB_SIZE`-byte and `SLAB_SIZE`-aligned mapping. The zero fill + /// doubles as the initial state of the carve counter and the link table. + pub unsafe fn init_header(base: NonNull, slab_idx: u32) { + // SAFETY: caller contract — the fresh exclusive mapping is aligned + // (SLAB_SIZE-aligned, far beyond SlabHeader's needs) and large + // enough for the header. + unsafe { + (*base.cast::().as_ptr()).slab_idx = slab_idx; + } + } + + const fn header(&self) -> &SlabHeader { + // SAFETY: type invariant — the mapping is live, SLAB_SIZE-aligned + // (far beyond SlabHeader's needs), and its header was initialized + // before publication; the only non-atomic field is never written + // again afterwards. + unsafe { self.base.cast::().as_ref() } + } + + /// This slab's index in its class's slab table. + pub const fn slab_idx(&self) -> usize { + self.header().slab_idx as usize + } + + /// The monotonic carve counter. + pub const fn carved(&self) -> &AtomicU64 { + &self.header().carved + } + + const fn blocks(&self) -> usize { + BLOCKS_PER_SLAB[self.class] + } + + /// The free-list link slot of `block_idx`. + pub fn link(&self, block_idx: usize) -> &AtomicU32 { + debug_assert!(block_idx < self.blocks()); + #[expect( + clippy::cast_ptr_alignment, + reason = "the link table starts at offset 16 of a SLAB_SIZE-aligned slab, so entries are 4-byte aligned" + )] + // SAFETY: type invariant plus the bound above put the slot within + // the slab's link table; slots are only ever accessed atomically, + // and the mapping outlives any borrow. + unsafe { + AtomicU32::from_ptr( + self.base.as_ptr().add(size_of::()).cast::().add(block_idx), + ) + } + } + + /// The address of block `block_idx`. + pub fn block(&self, block_idx: usize) -> NonNull { + debug_assert!(block_idx < self.blocks()); + let offset = block_area_offset(self.blocks()) + block_idx * block_size(self.class); + // SAFETY: type invariant plus the bound above keep `offset` within + // the SLAB_SIZE mapping. + unsafe { self.base.add(offset) } + } + + /// The index of a block previously returned by [`Slab::block`]. + pub fn block_index(&self, block: NonNull) -> usize { + let offset = + block.as_ptr().addr() - self.base.as_ptr().addr() - block_area_offset(self.blocks()); + debug_assert_eq!(offset % block_size(self.class), 0); + let block_idx = offset / block_size(self.class); + debug_assert!(block_idx < self.blocks()); + block_idx + } +} diff --git a/crates/fspy_alloc/src/sys.rs b/crates/fspy_alloc/src/sys.rs new file mode 100644 index 000000000..4fecdc085 --- /dev/null +++ b/crates/fspy_alloc/src/sys.rs @@ -0,0 +1,25 @@ +//! The backing-memory provider abstraction. +//! +//! The pool is generic over [`Sys`] so tests (and Miri) can substitute a +//! provider backed by the host allocator, while the real allocator obtains +//! anonymous mappings from the kernel (see the `mmap` module). + +use core::ptr::NonNull; + +/// Provides zero-initialized, aligned memory regions. +/// +/// Implementations must themselves be async-signal-safe and fork-safe: no +/// locks, no thread-locals, no libc malloc. +pub trait Sys { + /// Maps `size` bytes of zero-initialized memory aligned to `align` + /// (a power of two). Returns `None` when memory is exhausted. + fn map(size: usize, align: usize) -> Option>; + + /// Releases a region previously returned by [`Sys::map`]. + /// + /// # Safety + /// + /// `ptr` must have been returned by `Sys::map(size, align)` with the same + /// `size` and `align`, and must not be accessed afterwards. + unsafe fn unmap(ptr: NonNull, size: usize, align: usize); +} diff --git a/crates/fspy_alloc/tests/global_alloc.rs b/crates/fspy_alloc/tests/global_alloc.rs new file mode 100644 index 000000000..ab0670ae5 --- /dev/null +++ b/crates/fspy_alloc/tests/global_alloc.rs @@ -0,0 +1,84 @@ +//! Installs [`FspyAlloc`] as this test binary's global allocator, so every +//! allocation — including the test harness's own — exercises the allocator +//! end-to-end over real memory mappings. +#![cfg(all(unix, not(miri)))] // Miri covers the pool core via its mockable backend instead. + +use std::{collections::BTreeMap, sync::mpsc, thread}; + +use fspy_alloc::FspyAlloc; + +#[global_allocator] +static GLOBAL: FspyAlloc = FspyAlloc::new(); + +#[test] +fn collections_round_trip() { + let mut map = BTreeMap::new(); + for i in 0..1000_u32 { + let len = usize::try_from(i % 300).unwrap(); + map.insert(i, vec![0xAB_u8; len]); + } + assert_eq!(map.len(), 1000); + for (i, bytes) in &map { + assert_eq!(bytes.len(), usize::try_from(i % 300).unwrap()); + assert!(bytes.iter().all(|byte| *byte == 0xAB)); + } +} + +#[test] +fn vec_growth_reallocs_preserve_contents() { + let mut bytes = Vec::new(); + for i in 0..1_000_000_usize { + bytes.push(u8::try_from(i % 251).unwrap()); + } + for (i, byte) in bytes.iter().enumerate() { + assert_eq!(usize::from(*byte), i % 251); + } +} + +#[test] +fn large_and_zeroed_allocations() { + // Direct-mapped (beyond the largest size class), via the zeroing path. + let large = vec![0_u8; 5 * 1024 * 1024]; + assert!(large.iter().all(|byte| *byte == 0)); + // And through the plain path. + let boxed: Box<[u8]> = vec![7_u8; 300 * 1024].into_boxed_slice(); + assert!(boxed.iter().all(|byte| *byte == 7)); +} + +#[test] +fn many_small_allocations_span_slabs() { + // More live 16-byte-class blocks than one slab holds, forcing the pool + // through several slab installations (and the holding Vec through the + // direct-mapped path as it grows). + let boxes: Vec> = (0..120_000).map(|_| Box::new([0xEE_u8; 8])).collect(); + assert!(boxes.iter().all(|block| block.iter().all(|byte| *byte == 0xEE))); +} + +#[test] +fn threaded_producers_and_consumers() { + let threads = 8; + let per_thread = 5_000_usize; + let (sender, receiver) = mpsc::channel::>(); + thread::scope(|scope| { + for t in 0..threads { + let sender = sender.clone(); + scope.spawn(move || { + for i in 0..per_thread { + // Vary sizes across classes; blocks are freed (and often + // allocated) on the consumer thread below. + let len = 1 + (i * 37 + t * 101) % 5000; + sender.send(vec![u8::try_from(t % 251).unwrap(); len]).unwrap(); + } + }); + } + drop(sender); + let mut message_count = 0_usize; + for bytes in receiver { + message_count += 1; + assert!(!bytes.is_empty()); + let first = bytes[0]; + assert!(bytes.iter().all(|byte| *byte == first)); + } + assert_eq!(message_count, threads * per_thread); + }); +} diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 4b89bbd85..70a554563 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -12,6 +12,7 @@ anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } ctor = { workspace = true } +fspy_alloc = { workspace = true } fspy_shared = { workspace = true } fspy_shared_unix = { workspace = true } libc = { workspace = true } diff --git a/crates/fspy_preload_unix/src/lib.rs b/crates/fspy_preload_unix/src/lib.rs index 6c4e10b3e..a9e667970 100644 --- a/crates/fspy_preload_unix/src/lib.rs +++ b/crates/fspy_preload_unix/src/lib.rs @@ -1,6 +1,16 @@ // Compile as an empty crate on non-unix targets and on musl (where seccomp // alone handles access tracking). +/// This library interposes libc functions that POSIX declares +/// async-signal-safe (`open`, `stat`, `execve`, ...), so its own code may run +/// inside signal handlers and in the child of `fork()` in a multithreaded +/// process — contexts where taking the libc allocator's locks can deadlock. +/// Route every Rust allocation in this cdylib through fspy's lock-free, +/// mmap-backed allocator instead. +#[cfg(all(unix, not(target_env = "musl")))] +#[global_allocator] +static GLOBAL_ALLOCATOR: fspy_alloc::FspyAlloc = fspy_alloc::FspyAlloc::new(); + #[cfg(all(unix, not(target_env = "musl")))] mod client; #[cfg(all(unix, not(target_env = "musl")))] From 87be7e0147587f05424affcd50ce2dbf776a4302 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 08:05:46 +0800 Subject: [PATCH 2/4] feat(fspy): replace the pool allocator with per-call arenas over a chunk pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rework fspy_alloc from a size-class global allocator into three layers: MmapAllocator (stateless, every allocation is a fresh kernel mapping), ChunkPool (a lock-free cache of 64 KiB chunks), and arena(), the only public entry, which hands each intercepted call its own bump arena (bump_scope::Bump) drawing chunks from the process-wide pool. Use the arena for the first preload call site: joining a directory and a relative path when resolving fd-relative opens. Add an access-relative benchmark suite so this lane — working-directory resolution plus path joining — is measured; the existing absolute-path suite never enters it. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 12 + Cargo.toml | 4 +- crates/fspy_alloc/Cargo.toml | 7 + crates/fspy_alloc/src/class.rs | 38 - crates/fspy_alloc/src/lib.rs | 223 +++-- crates/fspy_alloc/src/mapping.rs | 62 -- crates/fspy_alloc/src/mmap.rs | 310 +++--- crates/fspy_alloc/src/pool.rs | 914 ++++++------------ crates/fspy_alloc/src/slab.rs | 203 ---- crates/fspy_alloc/src/sys.rs | 25 - crates/fspy_alloc/tests/global_alloc.rs | 84 -- crates/fspy_preload_unix/Cargo.toml | 1 + .../fspy_preload_unix/src/client/convert.rs | 22 +- crates/fspy_preload_unix/src/lib.rs | 10 - 14 files changed, 580 insertions(+), 1335 deletions(-) delete mode 100644 crates/fspy_alloc/src/class.rs delete mode 100644 crates/fspy_alloc/src/mapping.rs delete mode 100644 crates/fspy_alloc/src/slab.rs delete mode 100644 crates/fspy_alloc/src/sys.rs delete mode 100644 crates/fspy_alloc/tests/global_alloc.rs diff --git a/Cargo.lock b/Cargo.lock index 6182d52d4..b783495c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -349,6 +349,15 @@ dependencies = [ "serde", ] +[[package]] +name = "bump-scope" +version = "2.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4747cadbe8ffeff2a58cb86bcb35f9ec490f2329e12a88a696bccfd3dc97bf4d" +dependencies = [ + "allocator-api2", +] + [[package]] name = "bumpalo" version = "3.19.1" @@ -1247,6 +1256,8 @@ dependencies = [ name = "fspy_alloc" version = "0.0.0" dependencies = [ + "allocator-api2", + "bump-scope", "rustix", ] @@ -1302,6 +1313,7 @@ dependencies = [ name = "fspy_preload_unix" version = "0.0.0" dependencies = [ + "allocator-api2", "anyhow", "artifact_profile", "bstr", diff --git a/Cargo.toml b/Cargo.toml index 4385383c4..f1c766099 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -39,6 +39,7 @@ multiple_crate_versions = "allow" future_not_send = "allow" [workspace.dependencies] +allocator-api2 = { version = "0.2", default-features = false } artifact_profile = { path = "crates/artifact_profile" } anstream = "1.0.0" anyhow = "1.0.103" @@ -51,6 +52,7 @@ bindgen = "0.72.1" bitflags = "2.10.0" brush-parser = "0.4.0" bstr = { version = "1.12.0", default-features = false, features = ["alloc", "std"] } +bump-scope = { version = "2", default-features = false, features = ["allocator-api2-02"] } bumpalo = { version = "3.17.0", features = ["collections"] } bytemuck = { version = "1.23.0", features = ["extern_crate_alloc", "must_cast"] } cc = "1.2.39" @@ -121,7 +123,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" -rustix = { version = "1", default-features = false, features = ["mm", "param"] } +rustix = { version = "1", default-features = false, features = ["mm", "param", "use-libc-auxv"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml index 809d0f7e8..edcb89eaa 100644 --- a/crates/fspy_alloc/Cargo.toml +++ b/crates/fspy_alloc/Cargo.toml @@ -8,7 +8,14 @@ publish = false doctest = false [target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true } +bump-scope = { workspace = true } rustix = { workspace = true } +[target.'cfg(unix)'.dev-dependencies] +# The `alloc` feature provides `Global`, letting tests run the pool against +# the host allocator (and thus under Miri). +allocator-api2 = { workspace = true, features = ["alloc"] } + [lints] workspace = true diff --git a/crates/fspy_alloc/src/class.rs b/crates/fspy_alloc/src/class.rs deleted file mode 100644 index df8d3797a..000000000 --- a/crates/fspy_alloc/src/class.rs +++ /dev/null @@ -1,38 +0,0 @@ -//! Size-class policy: which layouts the pool serves, and at what block size. - -use core::alloc::Layout; - -const MIN_CLASS_SHIFT: u32 = 4; -const MAX_CLASS_SHIFT: u32 = 16; -pub const CLASS_COUNT: usize = (MAX_CLASS_SHIFT - MIN_CLASS_SHIFT) as usize + 1; -const MIN_BLOCK_SIZE: usize = 1 << MIN_CLASS_SHIFT; -pub const MAX_BLOCK_SIZE: usize = 1 << MAX_CLASS_SHIFT; -/// Block areas start at this alignment within a slab, making it the largest -/// alignment the pool can serve; stricter layouts map directly. -pub const MAX_POOL_ALIGN: usize = 4096; - -pub const fn block_size(class: usize) -> usize { - 1 << (MIN_CLASS_SHIFT as usize + class) -} - -/// Returns the size class for `layout`, or `None` if the request must be -/// mapped directly (too large or over-aligned). -pub const fn class_of(layout: Layout) -> Option { - if layout.align() > MAX_POOL_ALIGN { - return None; - } - let mut size = layout.size(); - // A block of `size >= align` at a `min(block size, 4 KiB)` boundary is - // aligned to `align` (both are powers of two and `align <= 4 KiB`). - if size < layout.align() { - size = layout.align(); - } - if size < MIN_BLOCK_SIZE { - size = MIN_BLOCK_SIZE; - } - if size > MAX_BLOCK_SIZE { - return None; - } - let shift = size.next_power_of_two().trailing_zeros(); - Some((shift - MIN_CLASS_SHIFT) as usize) -} diff --git a/crates/fspy_alloc/src/lib.rs b/crates/fspy_alloc/src/lib.rs index 90b2524cf..24b97020e 100644 --- a/crates/fspy_alloc/src/lib.rs +++ b/crates/fspy_alloc/src/lib.rs @@ -1,124 +1,147 @@ -//! Lock-free, async-signal-safe global allocator for the fspy preload library. +//! Async-signal-safe allocation for the fspy preload library. //! //! The preload library interposes libc functions that POSIX declares -//! async-signal-safe (`open`, `stat`, `execve`, ...). Programs may call these -//! from signal handlers, and — more commonly — from the child of `fork()` in a -//! multithreaded process, where only async-signal-safe calls are permitted: -//! the libc allocator's locks may be held forever by threads that no longer -//! exist after the fork. Routing the preload's Rust allocations through this -//! allocator keeps them safe in both contexts: +//! async-signal-safe (`open`, `stat`, `execve`, ...). Its code therefore runs +//! inside signal handlers and in the child of `fork()` in a multithreaded +//! process — contexts where taking libc malloc's locks can deadlock on a +//! suspended or vanished lock holder, so the preload's own allocations must +//! never go through libc malloc. //! -//! - **No locks.** Every state transition is a lock-free compare-and-swap -//! loop: an attempt only retries because another running thread completed -//! its operation, so nothing ever waits on state that a thread which -//! vanished at `fork()` — or sits suspended under a signal handler — would -//! have to release. (Lock-free, not wait-free: an individual operation has -//! no fixed retry bound under active contention.) -//! - **No thread-locals.** TLS first-touch allocates through libc malloc on -//! some platforms (macOS thread-local variables), which would reintroduce -//! the hazard this crate exists to remove. -//! - **mmap-backed.** Memory comes straight from the kernel. On Linux the -//! allocator relies on nothing from libc: mapping syscalls are issued -//! directly (rustix's raw backend) and even the page size is discovered by -//! probing with raw syscalls. On macOS, which has no stable raw-syscall -//! ABI, calls go through the thin libSystem stubs. libc malloc is never -//! called anywhere. -//! -//! Design: power-of-two size classes (16 B ..= 64 KiB) carve blocks out of -//! 1 MiB slabs; freed blocks recycle through a per-class Treiber free list -//! made ABA-safe by a generation tag. Requests larger than the biggest class -//! (or over-aligned beyond 4 KiB) map and unmap directly. See the `pool` -//! module for the details. -//! -//! Because the allocator is a `const`-initialized static with no lazy setup, -//! it works from the very first allocation in the process — even before the -//! preload library's constructor runs. +//! This crate stacks three layers, and exposes only the top one, [`arena`]. +//! `MmapAllocator` is the bottom: a stateless allocator where every +//! allocation is a fresh anonymous mapping from the kernel. `ChunkPool` +//! sits on top of it and caches fixed-size chunks, so that frequent short +//! tracing calls can reuse memory instead of paying two syscalls per call. +//! [`arena`] creates one `bump_scope::Bump` per intercepted call, drawing +//! its chunks from the process-wide pool and returning them on drop. +// Compile as an empty crate on non-unix targets: the allocator backs the +// unix preload library. +#![cfg(unix)] #![cfg_attr(not(test), no_std)] -// Compile as an empty crate on non-unix targets: the allocator backs the unix -// preload library. A Windows backend can be added alongside `sys::Mmap` if -// the Windows preload ever needs one. - -#[cfg(unix)] -mod class; -#[cfg(unix)] -mod mapping; -#[cfg(unix)] mod mmap; -#[cfg(unix)] mod pool; -#[cfg(unix)] -mod slab; -#[cfg(unix)] -mod sys; -#[cfg(unix)] -use core::{ - alloc::{GlobalAlloc, Layout}, - ptr::{self, NonNull}, +use allocator_api2::alloc::Allocator; +use bump_scope::{ + Bump, + alloc::compat::AllocatorApi2V02Compat, + settings::{BumpAllocatorSettings, BumpSettings}, }; +use mmap::MmapAllocator; +use pool::ChunkPool; -#[cfg(unix)] -use crate::{mmap::Mmap, pool::Pool}; +/// Every cached chunk is 64 KiB: a whole multiple of the page size on all +/// supported targets, and big enough that most intercepted calls fit their +/// allocations into a single chunk. +const CHUNK_SIZE: usize = 64 * 1024; +/// The alignment chunks are allocated with. Must be at least the alignment +/// `bump_scope::Bump` uses for its chunk requests — 16 (see +/// [`MmapAllocator`]'s docs for the links). bump-scope does not export that +/// constant, so the `bump_chunk_requests_fit_the_pool_gates` test pins the +/// fit instead: it fails if a bump-scope upgrade ever requests chunks the +/// pool would refuse. +const CHUNK_ALIGN: usize = 16; +/// At most this many chunks stay cached, capping retained memory at +/// `SLOTS * CHUNK_SIZE` = 4 MiB. +const SLOTS: usize = 64; -/// A lock-free, async-signal-safe, fork-safe [`GlobalAlloc`] implementation. -/// -/// Intended to be installed as the `#[global_allocator]` of the fspy preload -/// library. All memory comes from anonymous mappings; libc malloc is never -/// called, no locks are taken, and no thread-local state is used. -/// -/// Capacity is bounded by design: each size class can hold at most 256 slabs -/// of 1 MiB (roughly 200 MiB per class). Requests beyond that — far outside -/// anything the preload library does — fail like any other out-of-memory -/// condition (`alloc` returns null). -#[cfg(unix)] -pub struct FspyAlloc { - pool: Pool, -} +/// The process-wide chunk pool. `const`-initialized, so it works from the +/// first allocation on — even before any constructor has run. +static CHUNK_POOL: ChunkPool = ChunkPool::new(); -#[cfg(unix)] -impl FspyAlloc { - /// Creates the allocator. `const` so it can back a `static` with no - /// runtime initialization. - #[must_use] - pub const fn new() -> Self { - Self { pool: Pool::new() } +/// `Bump::unallocated` requires its base allocator to implement `Default` +/// (an arena without chunks has nowhere to store an allocator value, so it +/// conjures one on first use). Point defaulted references at the +/// process-wide pool. As an allocator, `&ChunkPool` already works through +/// allocator-api2's blanket `impl Allocator for &A`. +impl Default for &'static ChunkPool { + fn default() -> Self { + &CHUNK_POOL } } -#[cfg(unix)] -impl Default for FspyAlloc { - fn default() -> Self { - Self::new() - } +/// Creates a fresh bump arena for one intercepted call, backed by the +/// process-wide chunk pool. +/// +/// Creating the arena allocates nothing; the first allocation grabs a whole +/// chunk — usually a recycled one, so most calls touch no syscalls at all. +/// Deallocation only takes back the most recent allocation (bump-arena +/// semantics); everything is freed at once when the arena is dropped, and +/// its chunks go back to the pool. +/// +/// The arena itself is single-owner — use one per call, do not share it +/// across threads. Creating one is safe anywhere, any time: the pool +/// underneath works in signal handlers and in the child of `fork()` (see +/// `ChunkPool` and `MmapAllocator` in this crate's source for why). +#[must_use] +pub fn arena() -> impl Allocator { + // The default `Bump` settings, except the arena starts life without a + // chunk, so creating one allocates nothing. + type Settings = ::WithGuaranteedAllocated; + Bump::< + AllocatorApi2V02Compat<&'static ChunkPool>, + Settings, + >::unallocated() } -// SAFETY: `Pool` hands out blocks that are non-null, at least `layout.size()` -// bytes large, aligned to at least `layout.align()`, and exclusively owned -// until returned via `dealloc`. Allocation failure is reported as null, and -// none of the methods unwind. -#[cfg(unix)] -unsafe impl GlobalAlloc for FspyAlloc { - unsafe fn alloc(&self, layout: Layout) -> *mut u8 { - self.pool.alloc(layout).map_or(ptr::null_mut(), NonNull::as_ptr) - } +#[cfg(test)] +mod tests { + use core::alloc::Layout; - unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { - let Some(ptr) = NonNull::new(ptr) else { return }; - // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this - // allocator for this `layout`. - unsafe { self.pool.dealloc(ptr, layout) } - } + use allocator_api2::alloc::Global; - unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { - self.pool.alloc_zeroed(layout).map_or(ptr::null_mut(), NonNull::as_ptr) + use super::*; + + /// Pins the fit between `Bump`'s chunk requests and the pool's gates + /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]). + /// bump-scope keeps those request parameters private, so this test is + /// the enforcement: it fails if an upgrade ever changes them. + #[test] + fn bump_chunk_requests_fit_the_pool_gates() { + let pool = ChunkPool::::new_in(Global); + // `try_new_in` requests the first chunk right away; if that request + // were over-aligned or oversized, the pool would refuse and this + // would be an error. + let bump: Bump>> = + Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); + let block = bump.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); + let block_addr = block.cast::().as_ptr().addr(); + drop(bump); + + // The chunk went back into the pool's cache (proving it was served + // as a chunk, not passed through): the next chunk-sized request + // returns the block's surroundings. + let recycled = pool.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); + assert_eq!(recycled.len(), CHUNK_SIZE); + let base = recycled.cast::().as_ptr().addr(); + assert!((base..base + CHUNK_SIZE).contains(&block_addr)); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(recycled.cast(), Layout::from_size_align(100, 8).unwrap()) }; } - unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { - let Some(ptr) = NonNull::new(ptr) else { return ptr::null_mut() }; - // SAFETY: per the GlobalAlloc contract, `ptr` was returned by this - // allocator for this `layout`, and `new_size` is non-zero. - unsafe { self.pool.realloc(ptr, layout, new_size) }.map_or(ptr::null_mut(), NonNull::as_ptr) + #[test] + #[cfg(not(miri))] + fn arena_allocates_and_returns_chunks_to_the_pool() { + let layout = Layout::from_size_align(100, 8).unwrap(); + + let first_arena = arena(); + let first = first_arena.allocate(layout).unwrap(); + assert!(first.len() >= 100); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { first.cast::().as_ptr().write_bytes(0x5A, 100) }; + let second = first_arena.allocate(layout).unwrap(); + assert_ne!(first.cast::().as_ptr().addr(), second.cast::().as_ptr().addr()); + let first_addr = first.cast::().as_ptr().addr(); + // Everything dies at once; the chunk goes back to the pool. + drop(first_arena); + + // No other test touches the process-wide pool, so a new arena draws + // the same chunk back and its first allocation lands at the same + // address. + let second_arena = arena(); + let again = second_arena.allocate(layout).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), first_addr); } } diff --git a/crates/fspy_alloc/src/mapping.rs b/crates/fspy_alloc/src/mapping.rs deleted file mode 100644 index d3f53e726..000000000 --- a/crates/fspy_alloc/src/mapping.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! Owned memory regions obtained from a [`Sys`] provider. - -use core::{marker::PhantomData, mem, ptr::NonNull}; - -use crate::sys::Sys; - -/// An owned region obtained from `S`, released on drop. -/// -/// This is the only place that calls [`Sys::unmap`]: pool code either lets a -/// `Mapping` drop (probe scratch, install races, freed large allocations) or -/// deliberately leaks it with [`Mapping::into_raw`] (published slabs, live -/// large allocations). Reconstructing ownership from a raw pointer via -/// [`Mapping::from_raw`] is the single unsafe step. -pub struct Mapping { - ptr: NonNull, - size: usize, - align: usize, - sys: PhantomData S>, -} - -impl Mapping { - /// Maps `size` bytes of zero-initialized memory aligned to `align` - /// (a power of two). Returns `None` when memory is exhausted. - pub fn new(size: usize, align: usize) -> Option { - let ptr = S::map(size, align)?; - Some(Self { ptr, size, align, sys: PhantomData }) - } - - /// Reclaims ownership of a mapping previously released with - /// [`Mapping::into_raw`]. - /// - /// # Safety - /// - /// `ptr` must have come from `Mapping::::into_raw` (or `Sys::map`) - /// with exactly this `size` and `align`, the region must not be in use, - /// and ownership must not be reclaimed twice. - pub unsafe fn from_raw(ptr: NonNull, size: usize, align: usize) -> Self { - Self { ptr, size, align, sys: PhantomData } - } - - /// The mapped region's base address. - pub const fn ptr(&self) -> NonNull { - self.ptr - } - - /// Releases ownership without unmapping; the region lives until (unless) - /// [`Mapping::from_raw`] reclaims it. - pub const fn into_raw(self) -> NonNull { - let ptr = self.ptr; - mem::forget(self); - ptr - } -} - -impl Drop for Mapping { - fn drop(&mut self) { - // SAFETY: this type owns the mapping (constructed from `Sys::map` - // directly or via the `from_raw` contract), and after drop nothing - // can use it. - unsafe { S::unmap(self.ptr, self.size, self.align) } - } -} diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/fspy_alloc/src/mmap.rs index 053356b56..3cc09d876 100644 --- a/crates/fspy_alloc/src/mmap.rs +++ b/crates/fspy_alloc/src/mmap.rs @@ -1,211 +1,161 @@ -//! Kernel-backed [`Sys`] provider: anonymous mappings via direct syscalls. -//! -//! On Linux the provider relies on nothing from libc — rustix issues raw -//! syscalls, and even the page size is discovered with raw syscalls (see -//! [`page_size`]). macOS has no stable raw-syscall ABI, so calls go through -//! the thin libSystem stubs there. +//! Page-granularity allocator backed by anonymous memory mappings. use core::{ + alloc::Layout, ptr::{self, NonNull}, - sync::atomic::{AtomicUsize, Ordering}, }; +use allocator_api2::alloc::{AllocError, Allocator}; use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; -use crate::sys::Sys; - -/// Returns the kernel page size, or `None` when it cannot be determined. +/// A stateless allocator: every allocation is a fresh anonymous mapping and +/// every deallocation an `munmap`. /// -/// On Linux the value is discovered with raw syscalls only — no libc, no -/// `/proc`, no minimum kernel version (see [`query_page_size`]). On other -/// unix platforms (macOS, where every syscall goes through libSystem by -/// platform contract anyway) it comes from `sysconf` via rustix. Either -/// way the result is a process constant, validated as a power of two and -/// cached, so `map` and `unmap` can never disagree on rounding; an -/// undeterminable page size fails the allocation rather than guessing. -fn page_size() -> Option { - static CACHE: AtomicUsize = AtomicUsize::new(0); - let cached = CACHE.load(Ordering::Relaxed); - if cached != 0 { - return Some(cached); - } - let page = query_page_size()?; - if !page.is_power_of_two() { - return None; - } - // First store wins; every query returns the same value, so the cache - // only avoids repeated probing. - let _ = CACHE.compare_exchange(0, page, Ordering::Relaxed, Ordering::Relaxed); - Some(page) -} - -/// Discovers the page size by probing, using nothing but raw syscalls. +/// # Why it is safe in signal handlers and forked children /// -/// `mprotect` fails with `EINVAL` unless its address is a multiple of the -/// page size, so re-protecting a scratch mapping at increasing -/// power-of-two offsets identifies the page size as the first offset the -/// kernel accepts (for a power-of-two page size P, the first power of two -/// P divides is P itself). A handful of syscalls, once per process: -/// async-signal-safe, fork-safe, and independent of libc, `/proc` -/// availability, and kernel version. +/// It holds no state at all — no locks, no free lists, no thread-locals; +/// the kernel does all the bookkeeping. A signal or a `fork()` can never +/// catch it holding a lock or a half-written structure, because there is +/// nothing to hold. Mapping syscalls go straight to the kernel on Linux +/// (rustix's raw backend) and through the thin libSystem wrappers on macOS. +/// The page size comes from static process data (`getauxval` on Linux, +/// `sysconf` on macOS) — no locks or allocation there either. /// -/// Why not `rustix::param::page_size()` here (it's fine on macOS)? On -/// Linux the allocator must not rely on libc, which rules out rustix's -/// `use-libc-auxv` (`getauxval`) configuration — and rustix's libc-free -/// fallback is unusable *inside* a global allocator: +/// # What it accepts /// -/// - Its lazy init tries `prctl(PR_GET_AUXV)` (kernel 6.4+ only) and -/// otherwise reads `/proc/self/auxv`; with rustix's `alloc` feature -/// enabled, that read path heap-allocates (`Vec`) — through *this* -/// allocator, whose `map` is the caller waiting on the page size — -/// recursing unboundedly. And we cannot pin `alloc` off: Cargo feature -/// unification lets any other rustix user in the build graph enable it -/// for our copy. -/// - It panics on read errors, truncated auxv, or both sources being -/// unavailable, where this allocator requires failure to surface as -/// `None` (panic formatting itself allocates, re-entering the same -/// uninitialized path). +/// The only intended caller is `bump_scope::Bump`, and `Bump` only ever +/// asks its base allocator for chunks ([single call site][site]). Chunks +/// are never zero-sized (a chunk always contains its own header), and +/// their alignment is [`max(MIN_CHUNK_ALIGN, header alignment)`][chunk-align], +/// [where `MIN_CHUNK_ALIGN` is 16][mca] — so 16 bytes in practice. /// -/// The probe has neither problem: no allocation, no panic, no minimum -/// kernel, and its worst case is `None`. -#[cfg(target_os = "linux")] -fn query_page_size() -> Option { - use rustix::mm::{MprotectFlags, mprotect}; - - /// Smallest page size of any Linux configuration; the probe starts - /// here. - const MIN_PROBE_PAGE: usize = 4096; - /// Generous upper bound for the probe; no supported configuration - /// uses larger pages. - const MAX_PROBE_PAGE: usize = 1 << 20; +/// This allocator accepts more than that, because mapped memory gives the +/// extra range away for free: any non-zero size, and any alignment up to +/// the page size (mappings are always page-aligned). It refuses only two +/// kinds of request, which would each need extra code and never happen: +/// zero-sized layouts (they would need fake dangling blocks) and alignment +/// above the page size (it would need mapping extra space and trimming the +/// misaligned edges). The [`Allocator`] contract allows refusing any +/// request; refused requests get [`AllocError`]. +/// +/// # Cost +/// +/// Every allocation takes whole pages (4 KiB at least, 16 KiB on Apple +/// silicon) and one syscall, and so does every deallocation. This +/// allocator is meant to sit below a chunk pool and bump arenas, not to +/// serve small allocations directly. +/// +/// `allocate` returns the whole page-rounded block, and `bump_scope` [uses +/// the full returned length][fit], so none of the page is wasted. +/// +/// [mca]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/chunk/size_config.rs.html#8 +/// [chunk-align]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/chunk/size_config.rs.html#56-58 +/// [site]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/raw_bump.rs.html#865 +/// [fit]: https://docs.rs/bump-scope/2.3.3/src/bump_scope/raw_bump.rs.html#873-884 +#[derive(Clone, Copy, Debug, Default)] +pub struct MmapAllocator; + +pub fn page_size() -> usize { + // getauxval on Linux, sysconf on macOS: reads of static process data — + // no locks, no allocation, valid from the first instruction on. + rustix::param::page_size() +} - // Large enough that every probe below stays inside the mapping even - // after the kernel rounds the one-byte length up to a full page. - let scratch = map_anonymous(2 * MAX_PROBE_PAGE)?; - let mut page = None; - let mut offset = MIN_PROBE_PAGE; - while offset <= MAX_PROBE_PAGE { - // SAFETY: `scratch + offset` (plus the page-rounded single byte) - // lies within the scratch mapping, which we exclusively own; the - // protection flags match the mapping's existing ones. - let accepted = unsafe { - mprotect( - scratch.as_ptr().add(offset).cast(), - 1, - MprotectFlags::READ | MprotectFlags::WRITE, - ) +// SAFETY: returned blocks are non-null, page-aligned (at least +// `layout.align()` for every served layout), at least `layout.size()` bytes +// large (the returned length reports the exact mapped size), stay valid +// until deallocated, and distinct allocations never overlap. +unsafe impl Allocator for MmapAllocator { + fn allocate(&self, layout: Layout) -> Result, AllocError> { + let page = page_size(); + // Outside the served request profile (see the type docs): refuse + // rather than carry an over-alignment or dangling-block code path. + if layout.size() == 0 || layout.align() > page { + return Err(AllocError); } - .is_ok(); - if accepted { - page = Some(offset); - break; + // `checked_next_multiple_of` is total: `None` on overflow (already + // impossible — `Layout` caps sizes at `isize::MAX`) or a zero page + // size, with no power-of-two assumption to uphold. + let size = layout.size().checked_next_multiple_of(page).ok_or(AllocError)?; + // SAFETY: a fresh anonymous private mapping at no particular + // address has no memory-safety preconditions. + let ptr = unsafe { + mmap_anonymous( + ptr::null_mut(), + size, + ProtFlags::READ | ProtFlags::WRITE, + MapFlags::PRIVATE, + ) } - offset *= 2; + .map_err(|_| AllocError)?; + // Mapping results are page-aligned, which covers every served + // layout. + NonNull::new(ptr.cast::()) + .map(|ptr| NonNull::slice_from_raw_parts(ptr, size)) + .ok_or(AllocError) } - // SAFETY: releasing the scratch mapping created above. - let _ = unsafe { munmap(scratch.as_ptr().cast(), 2 * MAX_PROBE_PAGE) }; - page -} - -#[cfg(not(target_os = "linux"))] -#[expect( - clippy::unnecessary_wraps, - reason = "must match the signature of the fallible Linux probe variant" -)] -fn query_page_size() -> Option { - Some(rustix::param::page_size()) -} -const fn round_up(value: usize, align: usize) -> usize { - (value + align - 1) & !(align - 1) -} - -fn map_anonymous(len: usize) -> Option> { - // SAFETY: a fresh anonymous private mapping at no particular address - // has no memory-safety preconditions. - let ptr = unsafe { - mmap_anonymous(ptr::null_mut(), len, ProtFlags::READ | ProtFlags::WRITE, MapFlags::PRIVATE) - } - .ok()?; - NonNull::new(ptr.cast::()) -} - -/// Kernel-backed provider: anonymous mappings obtained through direct -/// syscalls, alignment achieved by over-mapping and trimming. -pub struct Mmap; - -impl Sys for Mmap { - fn map(size: usize, align: usize) -> Option> { - debug_assert!(align.is_power_of_two()); - let page = page_size()?; - let size = round_up(size.max(1), page); - - if align <= page { - // Mapping results are aligned to the (verified real) page - // size. - return map_anonymous(size); - } - - // Over-map by `align`, then unmap the misaligned head and the - // leftover tail. All cut points are page-aligned: `raw` and - // `aligned` are page-aligned, and `size`/`align` are multiples of - // the page size. - let raw = map_anonymous(size.checked_add(align)?)?; - let raw_addr = raw.as_ptr().addr(); - let aligned_addr = round_up(raw_addr, align); - let head = aligned_addr - raw_addr; - let tail = align - head; - if head > 0 { - // SAFETY: `[raw_addr, raw_addr + head)` lies within the fresh - // mapping and `raw_addr` is page-aligned. Failure is - // impossible for a region we own; if it happened anyway the - // pages would merely stay mapped. - let _ = unsafe { munmap(raw.as_ptr().cast(), head) }; - } - if tail > 0 { - let tail_start = raw.as_ptr().with_addr(aligned_addr + size); - // SAFETY: `[aligned_addr + size, raw_addr + size + align)` - // lies within the fresh mapping and its start is page-aligned - // (see above). Failure is impossible for a region we own. - let _ = unsafe { munmap(tail_start.cast(), tail) }; - } - // `aligned_addr` lies inside a successful mapping and so can - // never be zero, but checked construction costs nothing here. - core::num::NonZero::new(aligned_addr).map(|addr| raw.with_addr(addr)) + fn allocate_zeroed(&self, layout: Layout) -> Result, AllocError> { + // Fresh anonymous mappings are already zero-filled by the kernel. + self.allocate(layout) } - unsafe fn unmap(ptr: NonNull, size: usize, _align: usize) { - // A successful `map` proved the page size, so this cannot fail - // for a live mapping; if it somehow did, leaking the region is - // the only safe response. - let Some(page) = page_size() else { return }; - // Whether or not `map` trimmed for alignment, the retained region - // is exactly `[ptr, ptr + round_up(size, page))`. - let size = round_up(size.max(1), page); - // SAFETY: caller contract — `ptr`/`size` describe a live mapping - // returned by `map`. Failure is impossible for a region we own. + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + // The contract lets callers pass any size that fits the block, i.e. + // anything in `[requested, mapped]`; every such size rounds up to + // the mapped length. Zero-sized blocks are never allocated, so no + // dangling pointers can arrive here. Rounding cannot fail for a + // layout that fits a block we mapped; leaking the region is the + // safe response if it somehow did. + let Some(size) = layout.size().checked_next_multiple_of(page_size()) else { return }; + // SAFETY: caller contract — `ptr` was returned by `allocate` with a + // fitting layout and the block is no longer in use. Failure is + // impossible for a region we own. let _ = unsafe { munmap(ptr.as_ptr().cast(), size) }; } } #[cfg(all(test, not(miri)))] mod tests { - use super::Mmap; - use crate::sys::Sys; + use super::*; #[test] - fn real_mappings_are_aligned_zeroed_and_writable() { - for (size, align) in [(1, 1), (4096, 4096), (100, 1 << 20), (5 << 20, 4096)] { - let ptr = Mmap::map(size, align).unwrap(); - assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); - for i in 0..size { - // SAFETY: fresh exclusive mapping of at least `size` bytes. - assert_eq!(unsafe { ptr.as_ptr().add(i).read() }, 0); + fn blocks_are_aligned_zeroed_and_page_rounded() { + for (size, align) in [(1, 1), (100, 64), (4096, 4096), (5 << 20, 8)] { + let layout = Layout::from_size_align(size, align).unwrap(); + let block = MmapAllocator.allocate(layout).unwrap(); + assert!(block.len() >= size, "size {size}"); + assert_eq!(block.len() % page_size(), 0); + assert_eq!(block.cast::().as_ptr().addr() % align, 0, "align {align}"); + for i in 0..block.len() { + // SAFETY: fresh exclusive block of `block.len()` bytes. + assert_eq!(unsafe { block.cast::().as_ptr().add(i).read() }, 0); } - // SAFETY: fresh exclusive mapping of at least `size` bytes. - unsafe { ptr.as_ptr().write_bytes(0x5A, size) }; - // SAFETY: mapped above with the same size and alignment. - unsafe { Mmap::unmap(ptr, size, align) }; + // SAFETY: fresh exclusive block of at least `size` bytes. + unsafe { block.cast::().as_ptr().write_bytes(0x5A, size) }; + // SAFETY: allocated above; the layout fits the block. + unsafe { MmapAllocator.deallocate(block.cast(), layout) }; } } + + #[test] + fn deallocate_accepts_any_fitting_size() { + let requested = Layout::from_size_align(100, 8).unwrap(); + let block = MmapAllocator.allocate(requested).unwrap(); + // Deallocate with the *returned* size instead of the requested one — + // both are within the fit range the contract allows. + let fitting = Layout::from_size_align(block.len(), 8).unwrap(); + // SAFETY: allocated above; `fitting` is within the block's fit range. + unsafe { MmapAllocator.deallocate(block.cast(), fitting) }; + } + + #[test] + fn out_of_profile_requests_are_refused() { + // Zero-sized and over-page-aligned layouts are outside the served + // request profile and must fail cleanly, not misbehave. + let zero = Layout::from_size_align(0, 16).unwrap(); + assert!(MmapAllocator.allocate(zero).is_err()); + let over_aligned = Layout::from_size_align(64, 1 << 24).unwrap(); + assert!(MmapAllocator.allocate(over_aligned).is_err()); + } } diff --git a/crates/fspy_alloc/src/pool.rs b/crates/fspy_alloc/src/pool.rs index 6bb4f82f1..f0e3e4112 100644 --- a/crates/fspy_alloc/src/pool.rs +++ b/crates/fspy_alloc/src/pool.rs @@ -1,705 +1,363 @@ -//! Lock-free size-class pool. -//! -//! Layouts round up to a power-of-two class (see the `class` module) whose -//! blocks are carved from slabs (see the `slab` module for the memory -//! layout). -//! -//! # Concurrency -//! -//! Each class has a Treiber-stack free list plus a bump cursor over the -//! newest ("active") slab: -//! -//! - **alloc** pops the free list, or carves the next block off the active -//! slab, installing a fresh slab when the active one is exhausted. -//! - **dealloc** pushes the block back onto its class's free list. -//! -//! The 64-bit list head packs a 40-bit generation tag next to the 24-bit -//! block reference; the tag advances on every successful push and pop, which -//! makes the classic Treiber-stack ABA failure require a thread to stall -//! between its head load and CAS while other threads perform an exact -//! multiple of 2^40 head mutations. That is a probabilistic defense, not a -//! formal impossibility — see [`HEAD_TAG_SHIFT`]. Slab installation races -//! are resolved with a compare-and-swap on the slab table slot; the loser -//! simply returns its (never-published) mapping. -//! -//! The pool is **lock-free, not wait-free**: a CAS loop can retry -//! indefinitely under contention, but a retry only ever happens because -//! another thread completed an operation, so system-wide progress never -//! stalls — and, the property fork- and signal-safety actually require, no -//! operation ever waits on state that only a suspended or vanished thread -//! could release. Blocks are aligned to `min(block size, 4 KiB)`; requests -//! with stricter alignment (or size beyond the largest class) bypass the -//! pool and map directly. +//! A small cache of memory chunks between the bump arenas and the kernel. use core::{ alloc::Layout, - marker::PhantomData, - ptr, - ptr::NonNull, - sync::atomic::{AtomicPtr, AtomicU32, AtomicU64, Ordering}, + ptr::{self, NonNull}, + sync::atomic::{AtomicPtr, Ordering}, }; -use crate::{ - class::{CLASS_COUNT, class_of}, - mapping::Mapping, - slab::{BLOCKS_PER_SLAB, SLAB_SIZE, Slab}, - sys::Sys, -}; - -/// Bounded by the 8 bits reserved for slab indices in a packed block -/// reference. Caps each class at ~200 MiB. -const MAX_SLABS_PER_CLASS: usize = 256; - -/// Sentinel for "no block" in the 24-bit packed-reference field of a -/// free-list head. Never collides with a real reference: block indices stay -/// below `0xFFFF` (asserted with the slab geometry below). -const NO_BLOCK: u32 = 0x00FF_FFFF; -/// Sentinel for "no slab installed yet" in [`ClassState::active`]. -const NO_SLAB: u32 = u32::MAX; - -const _: () = { - assert!(MAX_SLABS_PER_CLASS <= 1 << 8, "slab index must fit in 8 bits"); - let mut class = 0; - while class < CLASS_COUNT { - assert!( - BLOCKS_PER_SLAB[class] <= 0xFFFF, - "block indices must fit in 16 bits and stay below the NO_BLOCK sentinel" - ); - class += 1; - } -}; - -/// Packs a block's location into 24 bits: slab index in bits 16..24, block -/// index in bits 0..16. -#[expect( - clippy::cast_possible_truncation, - reason = "callers pass indices bounded by MAX_SLABS_PER_CLASS and BLOCKS_PER_SLAB" -)] -const fn pack_ref(slab_idx: usize, block_idx: usize) -> u32 { - ((slab_idx as u32) << 16) | (block_idx as u32) -} - -const fn unpack_ref(packed: u32) -> (usize, usize) { - (((packed >> 16) & 0xFF) as usize, (packed & 0xFFFF) as usize) -} - -/// A free-list head is `[generation tag : 40 | packed block reference : 24]`. -/// The tag advances on every successful push and pop, so a stale -/// compare-and-swap can only succeed if its thread stalls between head load -/// and CAS while others perform an exact multiple of 2^40 head mutations — -/// not a formal impossibility, but hours of maximum-rate churn inside one -/// stalled instruction window. -const HEAD_TAG_SHIFT: u32 = 24; -const HEAD_TAG_MASK: u64 = (1 << 40) - 1; -const HEAD_REF_MASK: u64 = (1 << HEAD_TAG_SHIFT) - 1; - -/// Splits a free-list head into `(generation tag, packed block reference)`. -const fn head_parts(head: u64) -> (u64, u32) { - // The mask keeps the reference within 24 bits, so the cast is lossless. - (head >> HEAD_TAG_SHIFT, (head & HEAD_REF_MASK) as u32) -} - -const fn head_from_parts(tag: u64, block_ref: u32) -> u64 { - ((tag & HEAD_TAG_MASK) << HEAD_TAG_SHIFT) | block_ref as u64 -} - -/// Aligned to its own cache-line region so hot heads of different classes -/// don't false-share. -#[repr(align(128))] -struct ClassState { - /// Treiber free-list head: `[generation tag : 40 | packed ref : 24]`. - head: AtomicU64, - /// Index of the slab currently being carved, or [`NO_SLAB`] before the - /// first slab is installed. - active: AtomicU32, - /// Base addresses of installed slabs. Written once (null → mapping) and - /// never cleared. - slabs: [AtomicPtr; MAX_SLABS_PER_CLASS], -} - -impl ClassState { - const fn new() -> Self { - Self { - head: AtomicU64::new(head_from_parts(0, NO_BLOCK)), - active: AtomicU32::new(NO_SLAB), - slabs: [const { AtomicPtr::new(ptr::null_mut()) }; MAX_SLABS_PER_CLASS], - } - } -} - -/// A block handed out by [`Pool::alloc_block`], remembering whether it came -/// off the free list (may contain stale data) or was carved fresh off a slab -/// (still kernel-zeroed). -enum ClassBlock { - Recycled(NonNull), - Carved(NonNull), -} - -enum Carve { - Block(NonNull), - /// A new slab was (or concurrently got) installed; retry the allocation. - Retry, - /// Out of slab table entries or out of memory. - Exhausted, +use allocator_api2::alloc::{AllocError, Allocator}; + +use crate::MmapAllocator; + +/// A fixed-size cache of memory chunks. +/// +/// Sits between the bump arenas and the underlying allocator `A` — +/// [`MmapAllocator`] in real use, any allocator in tests — so the chunks +/// that a finished arena gives back reach the next arena without going back +/// to the kernel. +/// +/// The parameters are chosen by the layer above (the arena layer): every +/// cached chunk is exactly `CHUNK_SIZE` bytes and allocated with +/// `CHUNK_ALIGN`, which is also the strictest alignment the pool accepts +/// (`bump_scope::Bump` requests 16 — see [`MmapAllocator`]'s docs for the +/// links). The cache holds at most `SLOTS` chunks, so at most +/// `SLOTS * CHUNK_SIZE` bytes stay retained. +/// +/// Requests of up to `CHUNK_SIZE` bytes are served with a whole chunk — +/// a cached one, or a fresh one when the cache is empty. Bigger requests go +/// straight to the underlying allocator. Deallocation mirrors that: chunks +/// go back into the cache (or to the underlying allocator when it is full), +/// bigger blocks go to the underlying allocator directly. The size passed +/// to `deallocate` tells the two apart exactly: for a block served as a +/// chunk, every size the caller may legally pass is at most `CHUNK_SIZE`; +/// for a pass-through block, every legal size is bigger. +/// +/// Zero-sized requests and alignments above `CHUNK_ALIGN` are refused — +/// its only intended caller, `bump_scope::Bump`, never sends either. +/// +/// # Why it is safe in signal handlers and forked children +/// +/// The only state of the pool itself is a fixed array of atomic pointers, +/// one slot per cached chunk. Taking a chunk swaps a slot to null; +/// returning one swaps a null slot back. Every operation visits each slot +/// at most once and never waits — there are no locks to leave locked and no +/// lists to leave half-linked, so a thread that disappears mid-operation +/// (`fork`, a signal) can strand at most the one chunk it was holding, +/// never the pool. (The underlying allocator must give the same guarantee; +/// [`MmapAllocator`] does.) +/// +/// The pool is `const`-constructible, so it can live in a `static` and +/// work before anything else has run. +pub struct ChunkPool< + A: Allocator, + const CHUNK_SIZE: usize, + const CHUNK_ALIGN: usize, + const SLOTS: usize, +> { + slots: [AtomicPtr; SLOTS], + allocator: A, } -/// The allocator core, generic over its [`Sys`] memory provider. -pub struct Pool { - classes: [ClassState; CLASS_COUNT], - sys: PhantomData S>, -} - -impl Pool { - #[expect( - clippy::large_stack_arrays, - reason = "the class table (~30 KiB) is only ever materialized into a const-initialized static, never built on a runtime stack" - )] +impl + ChunkPool +{ + #[must_use] pub const fn new() -> Self { - Self { classes: [const { ClassState::new() }; CLASS_COUNT], sys: PhantomData } + Self::new_in(MmapAllocator) } +} - /// Allocates memory for `layout`. Returns `None` on exhaustion. - pub fn alloc(&self, layout: Layout) -> Option> { - match class_of(layout) { - Some(class) => match self.alloc_block(class)? { - ClassBlock::Recycled(ptr) | ClassBlock::Carved(ptr) => Some(ptr), - }, - // Deliberately leaked until `dealloc` reclaims ownership. - None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), +impl + ChunkPool +{ + pub const fn new_in(allocator: A) -> Self { + const { + assert!(CHUNK_SIZE > 0, "CHUNK_SIZE must not be zero"); + assert!(CHUNK_ALIGN.is_power_of_two(), "CHUNK_ALIGN must be a power of two"); } + Self { slots: [const { AtomicPtr::new(ptr::null_mut()) }; SLOTS], allocator } } - /// Like [`Pool::alloc`], but the returned memory is zeroed. - pub fn alloc_zeroed(&self, layout: Layout) -> Option> { - match class_of(layout) { - Some(class) => match self.alloc_block(class)? { - // Freshly carved blocks are still kernel-zeroed. - ClassBlock::Carved(ptr) => Some(ptr), - ClassBlock::Recycled(ptr) => { - // SAFETY: the block is freshly allocated, exclusively - // ours, and at least `layout.size()` bytes. - unsafe { ptr.as_ptr().write_bytes(0, layout.size()) }; - Some(ptr) - } - }, - // Fresh mappings are zeroed by the kernel; deliberately leaked - // until `dealloc` reclaims ownership. - None => Some(Mapping::::new(layout.size(), layout.align())?.into_raw()), - } + fn chunk_layout() -> Result { + Layout::from_size_align(CHUNK_SIZE, CHUNK_ALIGN).map_err(|_| AllocError) } +} - /// Releases memory obtained from this pool. - /// - /// # Safety - /// - /// `ptr` must have been returned by this pool for exactly this `layout` - /// and must not be used afterwards. - pub unsafe fn dealloc(&self, ptr: NonNull, layout: Layout) { - match class_of(layout) { - // SAFETY: caller contract — a `Some(class)` layout was served - // from the pool, so `ptr` is a live block of this class. - Some(class) => unsafe { self.push_free(ptr, class) }, - // SAFETY: caller contract — a `None` layout was served by the - // mapping path with these parameters and is no longer in use, so - // ownership can be reclaimed (and the region dropped). - None => drop(unsafe { Mapping::::from_raw(ptr, layout.size(), layout.align()) }), - } +impl + Default for ChunkPool +{ + fn default() -> Self { + Self::new_in(A::default()) } +} - /// Grows or shrinks an allocation, preserving contents up to the smaller - /// of the old and new sizes. Returns `None` on exhaustion (the original - /// allocation stays valid). - /// - /// # Safety - /// - /// `ptr` must have been returned by this pool for exactly `layout`, and - /// `new_size` must be non-zero. - pub unsafe fn realloc( - &self, - ptr: NonNull, - layout: Layout, - new_size: usize, - ) -> Option> { - let new_layout = Layout::from_size_align(new_size, layout.align()).ok()?; - let class = class_of(layout); - if class.is_some() && class == class_of(new_layout) { - // Same size class: the existing block already fits. - return Some(ptr); +// SAFETY: returned blocks are non-null, aligned to at least the accepted +// `layout.align()` (chunks are allocated with CHUNK_ALIGN, and stricter +// requests are refused), and at least `layout.size()` bytes large (chunks +// are exactly CHUNK_SIZE bytes and only serve requests up to that size; +// everything else comes from the underlying allocator with its own +// guarantees). A chunk leaves its slot before it is handed out and returns +// only on deallocation, so every block has exactly one owner and stays +// valid until deallocated. +unsafe impl + Allocator for ChunkPool +{ + fn allocate(&self, layout: Layout) -> Result, AllocError> { + if layout.size() == 0 || layout.align() > CHUNK_ALIGN { + return Err(AllocError); } - let new_ptr = self.alloc(new_layout)?; - // SAFETY: `new_ptr` is a fresh exclusive allocation of at least - // `new_size` bytes; `ptr` is valid for `layout.size()` bytes (caller - // contract); distinct allocations never overlap. - unsafe { - new_ptr.as_ptr().copy_from_nonoverlapping(ptr.as_ptr(), layout.size().min(new_size)); + if layout.size() > CHUNK_SIZE { + return self.allocator.allocate(layout); } - // SAFETY: caller contract — `ptr` came from this pool with `layout`. - unsafe { self.dealloc(ptr, layout) }; - Some(new_ptr) - } - - fn alloc_block(&self, class: usize) -> Option { - loop { - if let Some(ptr) = self.pop_free(class) { - return Some(ClassBlock::Recycled(ptr)); + for slot in &self.slots { + if slot.load(Ordering::Relaxed).is_null() { + continue; } - match self.carve(class) { - Carve::Block(ptr) => return Some(ClassBlock::Carved(ptr)), - Carve::Retry => {} - Carve::Exhausted => return None, + // Acquire pairs with the Release in `deallocate`; the swap makes + // this thread the chunk's only owner. A racing taker may empty + // the slot between the load and the swap — then the swap returns + // null and the scan moves on. + let chunk = slot.swap(ptr::null_mut(), Ordering::Acquire); + if let Some(chunk) = NonNull::new(chunk) { + return Ok(NonNull::slice_from_raw_parts(chunk, CHUNK_SIZE)); } } + // Cache empty: get a fresh chunk. + self.allocator.allocate(Self::chunk_layout()?) } - /// Returns the published slab of `class` at `slab_idx`, if installed. - fn published_slab(&self, class: usize, slab_idx: usize) -> Option { - let base = NonNull::new(self.classes[class].slabs[slab_idx].load(Ordering::Acquire))?; - // SAFETY: non-null slab-table entries are only published (with - // `Release`, paired with the `Acquire` above) after - // `Slab::init_header` ran on a fresh SLAB_SIZE-aligned mapping, and - // are never cleared or unmapped while the pool is in use. - Some(unsafe { Slab::from_published(base, class) }) - } - - /// Pops a block off the class's free list. - fn pop_free(&self, class: usize) -> Option> { - let state = &self.classes[class]; - loop { - let head = state.head.load(Ordering::Acquire); - let (tag, block_ref) = head_parts(head); - if block_ref == NO_BLOCK { - return None; - } - let (slab_idx, block_idx) = unpack_ref(block_ref); - // A listed block was carved from its slab, so the slab is always - // published; `None` here is unreachable in practice. - let slab = self.published_slab(class, slab_idx)?; - // Read the successor link before the CAS. If another thread pops - // this block first, the tag comparison below fails and the value - // read here is discarded; since links live in the atomic side - // table, the racing read itself is well-defined. - let next = slab.link(block_idx).load(Ordering::Relaxed); - let new_head = head_from_parts(tag.wrapping_add(1), next); - if state - .head - .compare_exchange_weak(head, new_head, Ordering::AcqRel, Ordering::Relaxed) - .is_ok() - { - return Some(slab.block(block_idx)); - } + unsafe fn deallocate(&self, ptr: NonNull, layout: Layout) { + if layout.size() > CHUNK_SIZE { + // A pass-through block: only requests bigger than CHUNK_SIZE get + // one, and all their legal deallocation sizes are bigger too. + // SAFETY: caller contract — the block came from `allocate` + // (which forwarded to the underlying allocator) with a fitting + // layout. + return unsafe { self.allocator.deallocate(ptr, layout) }; } - } - - /// Pushes a block onto the class's free list. - /// - /// # Safety - /// - /// `ptr` must be a block of `class` previously returned by this pool and - /// no longer in use. - unsafe fn push_free(&self, ptr: NonNull, class: usize) { - // SAFETY: caller contract — `ptr` is a live block of `class` from - // this pool. (The impossible `None` would merely leak the block.) - let Some(slab) = (unsafe { Slab::of_block(ptr, class) }) else { return }; - let block_idx = slab.block_index(ptr); - let packed = pack_ref(slab.slab_idx(), block_idx); - let link = slab.link(block_idx); - let state = &self.classes[class]; - loop { - let head = state.head.load(Ordering::Relaxed); - let (tag, head_ref) = head_parts(head); - link.store(head_ref, Ordering::Relaxed); - let new_head = head_from_parts(tag.wrapping_add(1), packed); - // `Release` publishes the link store above to the eventual popper. - if state - .head - .compare_exchange_weak(head, new_head, Ordering::Release, Ordering::Relaxed) + // A chunk: put it back into an empty slot, or give it back to the + // underlying allocator if the cache is full. Release pairs with the + // Acquire in `allocate`. A failed exchange means another thread just + // filled the slot; move on. + for slot in &self.slots { + if !slot.load(Ordering::Relaxed).is_null() { + continue; + } + if slot + .compare_exchange( + ptr::null_mut(), + ptr.as_ptr(), + Ordering::Release, + Ordering::Relaxed, + ) .is_ok() { return; } } + let Ok(chunk_layout) = Self::chunk_layout() else { return }; + // SAFETY: every block in this branch is a CHUNK_SIZE chunk that came + // from the underlying allocator with `chunk_layout` (caller contract + // plus the size routing above), and it no longer has an owner. + unsafe { self.allocator.deallocate(ptr, chunk_layout) }; } +} - /// Carves the next block off the class's active slab, installing a new - /// slab if the active one is exhausted (or none exists yet). - fn carve(&self, class: usize) -> Carve { - let state = &self.classes[class]; - let active = state.active.load(Ordering::Acquire); - // `active` is only published after its slab pointer, so a live - // `active` always resolves to a published slab. - if active != NO_SLAB - && let Some(slab) = self.published_slab(class, active as usize) - { - let carved_idx = slab.carved().fetch_add(1, Ordering::Relaxed); - if let Some(block_idx) = carved_to_block_idx(carved_idx, class) { - return Carve::Block(slab.block(block_idx)); - } - // Active slab exhausted; fall through to install the next one. - } - let next_idx = if active == NO_SLAB { 0 } else { active as usize + 1 }; - if next_idx >= MAX_SLABS_PER_CLASS { - return Carve::Exhausted; - } - self.install_slab(class, next_idx); - if self.published_slab(class, next_idx).is_none() { - // Our mapping failed and no other thread succeeded either. - return Carve::Exhausted; - } - #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] - let next_active = next_idx as u32; - // Advance `active`; losing the race just means another thread already - // advanced it. Either way the retry re-reads it. - let _ = - state.active.compare_exchange(active, next_active, Ordering::AcqRel, Ordering::Relaxed); - Carve::Retry - } - - /// Maps and publishes the slab at `slab_idx`, unless another thread beats - /// us to it (or the mapping fails, leaving the slot null). - fn install_slab(&self, class: usize, slab_idx: usize) { - let state = &self.classes[class]; - if !state.slabs[slab_idx].load(Ordering::Acquire).is_null() { - return; - } - let Some(mapping) = Mapping::::new(SLAB_SIZE, SLAB_SIZE) else { return }; - #[expect(clippy::cast_possible_truncation, reason = "bounded by MAX_SLABS_PER_CLASS")] - let idx = slab_idx as u32; - // SAFETY: `mapping` is a fresh, exclusive, zero-initialized, - // SLAB_SIZE-byte and SLAB_SIZE-aligned mapping. - unsafe { Slab::init_header(mapping.ptr(), idx) }; - if state.slabs[slab_idx] - .compare_exchange( - ptr::null_mut(), - mapping.ptr().as_ptr(), - Ordering::Release, - Ordering::Relaxed, - ) - .is_ok() - { - // Published: the slab now lives for the rest of the process. - let _ = mapping.into_raw(); - } - // Otherwise another thread installed this slot first; our - // never-published mapping is released when `mapping` drops. - } - - /// Tears down all slab mappings. Test-only: the global allocator lives in - /// a static and never releases its slabs, but tests (and Miri's leak - /// checker) want a clean shutdown. - #[cfg(test)] - fn unmap_all_slabs(&mut self) { - for state in &mut self.classes { - *state.head.get_mut() = head_from_parts(0, NO_BLOCK); - *state.active.get_mut() = NO_SLAB; - for slot in &mut state.slabs { - let slab = core::mem::replace(slot.get_mut(), ptr::null_mut()); - if let Some(base) = NonNull::new(slab) { - // SAFETY: `base` was leaked into the table by - // `install_slab` with these parameters; `&mut self` - // guarantees no concurrent (or future) use of its blocks. - drop(unsafe { Mapping::::from_raw(base, SLAB_SIZE, SLAB_SIZE) }); - } +impl Drop + for ChunkPool +{ + fn drop(&mut self) { + // Statics never drop; this runs for pools created in tests. + for slot in &mut self.slots { + let chunk = core::mem::replace(slot.get_mut(), ptr::null_mut()); + if let Some(chunk) = NonNull::new(chunk) { + let Ok(chunk_layout) = Self::chunk_layout() else { return }; + // SAFETY: cached chunks came from the underlying allocator + // with `chunk_layout`, and `&mut self` means no owner exists. + unsafe { self.allocator.deallocate(chunk, chunk_layout) }; } } } } -/// Converts a raw carve-counter value into a block index, or `None` if the -/// slab is exhausted. -fn carved_to_block_idx(carved: u64, class: usize) -> Option { - let idx = usize::try_from(carved).ok()?; - (idx < BLOCKS_PER_SLAB[class]).then_some(idx) -} - #[cfg(test)] mod tests { - use std::{boxed::Box, sync::mpsc, thread, vec::Vec}; - - use super::*; - use crate::class::{MAX_BLOCK_SIZE, MAX_POOL_ALIGN, block_size}; - - /// Host-allocator-backed [`Sys`] so the pool core runs under Miri and on - /// any platform. - struct TestSys; - - impl Sys for TestSys { - fn map(size: usize, align: usize) -> Option> { - let layout = Layout::from_size_align(size.max(1), align).ok()?; - // SAFETY: `layout` has non-zero size. - NonNull::new(unsafe { std::alloc::alloc_zeroed(layout) }) - } + use std::{thread, vec::Vec}; - unsafe fn unmap(ptr: NonNull, size: usize, align: usize) { - let layout = Layout::from_size_align(size.max(1), align).unwrap(); - // SAFETY: caller contract — `ptr` was returned by `map`, which - // used exactly this layout. - unsafe { std::alloc::dealloc(ptr.as_ptr(), layout) } - } - } + use allocator_api2::alloc::Global; - fn with_pool(test: impl FnOnce(&Pool)) { - let mut pool = Box::new(Pool::::new()); - test(&pool); - pool.unmap_all_slabs(); - } + use super::*; - fn layout(size: usize, align: usize) -> Layout { - Layout::from_size_align(size, align).unwrap() - } + const CHUNK_SIZE: usize = 64 * 1024; + const CHUNK_ALIGN: usize = 16; - #[test] - fn head_tag_wraps_within_its_field() { - let head = head_from_parts(HEAD_TAG_MASK, 42); - assert_eq!(head_parts(head), (HEAD_TAG_MASK, 42)); - // Advancing the maximal tag must wrap to zero without touching the - // reference bits. - let wrapped = head_from_parts(HEAD_TAG_MASK.wrapping_add(1), 42); - assert_eq!(head_parts(wrapped), (0, 42)); + fn pool() -> ChunkPool { + ChunkPool::new_in(Global) } - #[test] - fn packed_refs_cannot_collide_with_the_sentinel() { - for &blocks in &BLOCKS_PER_SLAB { - let max_ref = pack_ref(MAX_SLABS_PER_CLASS - 1, blocks - 1); - assert_ne!(max_ref, NO_BLOCK); - assert!(u64::from(max_ref) <= HEAD_REF_MASK); - } + fn layout(size: usize) -> Layout { + Layout::from_size_align(size, 8).unwrap() } #[test] - fn classify_boundaries() { - assert_eq!(class_of(layout(1, 1)), Some(0)); - assert_eq!(class_of(layout(16, 1)), Some(0)); - assert_eq!(class_of(layout(17, 1)), Some(1)); - assert_eq!(class_of(layout(MAX_BLOCK_SIZE, 8)), Some(CLASS_COUNT - 1)); - assert_eq!(class_of(layout(MAX_BLOCK_SIZE + 1, 8)), None); - // Alignment can raise the class. - assert_eq!(class_of(layout(8, 1024)), class_of(layout(1024, 8))); - // Over-aligned layouts bypass the pool. - assert_eq!(class_of(layout(16, MAX_POOL_ALIGN * 2)), None); + fn small_requests_get_whole_recycled_chunks() { + let pool = pool(); + let first = pool.allocate(layout(100)).unwrap(); + assert_eq!(first.len(), CHUNK_SIZE); + let first_addr = first.cast::().as_ptr().addr(); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { first.cast::().as_ptr().write_bytes(0xAB, 100) }; + // SAFETY: allocated above; the full returned size is a legal size to + // pass back. + unsafe { pool.deallocate(first.cast(), layout(CHUNK_SIZE)) }; + + // The chunk is cached now; the next small request must get it back. + let second = pool.allocate(layout(5000)).unwrap(); + assert_eq!(second.cast::().as_ptr().addr(), first_addr); + assert_eq!(second.len(), CHUNK_SIZE); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(second.cast(), layout(5000)) }; } #[test] - fn round_trip_every_class() { - with_pool(|pool| { - for class in 0..CLASS_COUNT { - let l = layout(block_size(class), 8); - let ptr = pool.alloc(l).unwrap(); - // SAFETY: fresh exclusive allocation of `block_size` bytes. - unsafe { ptr.as_ptr().write_bytes(0xAB, l.size()) }; - // SAFETY: reading back the block we just wrote. - let last = unsafe { ptr.as_ptr().add(l.size() - 1).read() }; - assert_eq!(last, 0xAB); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn deallocating_with_the_requested_size_also_recycles() { + let pool = pool(); + let block = pool.allocate(layout(100)).unwrap(); + let addr = block.cast::().as_ptr().addr(); + // SAFETY: allocated above; the requested size is a legal size to + // pass back, and must still be recognized as a chunk. + unsafe { pool.deallocate(block.cast(), layout(100)) }; + let again = pool.allocate(layout(100)).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), addr); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(again.cast(), layout(100)) }; } #[test] - fn block_alignment() { - with_pool(|pool| { - for align in [8, 64, 1024, MAX_POOL_ALIGN] { - let l = layout(24, align); - let ptr = pool.alloc(l).unwrap(); - assert_eq!(ptr.as_ptr().addr() % align, 0, "align {align}"); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn big_requests_bypass_the_cache() { + let pool = pool(); + let big = pool.allocate(layout(CHUNK_SIZE + 1)).unwrap(); + assert!(big.len() > CHUNK_SIZE); + // SAFETY: fresh exclusive block of at least CHUNK_SIZE + 1 + // bytes. + unsafe { big.cast::().as_ptr().write_bytes(0x5A, CHUNK_SIZE + 1) }; + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(big.cast(), layout(CHUNK_SIZE + 1)) }; + + // A small request afterwards gets a chunk, not the big block. + let small = pool.allocate(layout(100)).unwrap(); + assert_eq!(small.len(), CHUNK_SIZE); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(small.cast(), layout(100)) }; } #[test] - fn free_list_recycles_lifo() { - with_pool(|pool| { - let l = layout(100, 8); - let first = pool.alloc(l).unwrap(); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(first, l) }; - let second = pool.alloc(l).unwrap(); - assert_eq!(first, second); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(second, l) }; - }); + fn refuses_zero_size_and_over_chunk_alignment() { + let pool = pool(); + let zero = Layout::from_size_align(0, 16).unwrap(); + assert!(pool.allocate(zero).is_err()); + let over_aligned = Layout::from_size_align(64, CHUNK_ALIGN * 2).unwrap(); + assert!(pool.allocate(over_aligned).is_err()); } #[test] - fn carves_across_multiple_slabs() { - with_pool(|pool| { - // The largest class has the fewest blocks per slab, so a couple - // dozen live allocations force several slab installations. - let l = layout(MAX_BLOCK_SIZE, 8); - let count = BLOCKS_PER_SLAB[CLASS_COUNT - 1] * 3 + 1; - let blocks: Vec> = (0..count).map(|_| pool.alloc(l).unwrap()).collect(); - for (i, ptr) in blocks.iter().enumerate() { - assert!(blocks[..i].iter().all(|other| other != ptr), "duplicate block"); - // SAFETY: live exclusive allocation of `MAX_BLOCK_SIZE` bytes. - unsafe { ptr.as_ptr().write_bytes(0x5A, l.size()) }; - } - for ptr in blocks { - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; - } - }); + fn allocate_zeroed_scrubs_recycled_chunks() { + let pool = pool(); + let block = pool.allocate(layout(256)).unwrap(); + // SAFETY: fresh exclusive block of at least 256 bytes. + unsafe { block.cast::().as_ptr().write_bytes(0xFF, 256) }; + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(256)) }; + + // Recycled chunks are dirty, so the default `allocate_zeroed` (which + // we deliberately do not override here) must scrub them. + let zeroed = pool.allocate_zeroed(layout(256)).unwrap(); + for i in 0..256 { + // SAFETY: fresh exclusive block of at least 256 bytes. + assert_eq!(unsafe { zeroed.cast::().as_ptr().add(i).read() }, 0); + } + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(zeroed.cast(), layout(256)) }; } #[test] - fn large_allocations_bypass_pool() { - with_pool(|pool| { - for l in [layout(MAX_BLOCK_SIZE + 1, 8), layout(5 << 20, 8), layout(64, 8192)] { - let ptr = pool.alloc(l).unwrap(); - assert_eq!(ptr.as_ptr().addr() % l.align(), 0); - // SAFETY: fresh exclusive allocation of `l.size()` bytes. - unsafe { ptr.as_ptr().write_bytes(0xCD, l.size()) }; - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(ptr, l) }; + fn overflowing_the_cache_is_fine() { + // A tiny 4-slot pool so the full-cache path is genuinely exercised. + let pool = ChunkPool::::new_in(Global); + for _ in 0..2 { + let blocks: Vec<_> = (0..8).map(|_| pool.allocate(layout(1000)).unwrap()).collect(); + for block in blocks { + // SAFETY: allocated above with a fitting layout. The first 4 + // chunks fill the cache; the rest go back to the underlying + // allocator. + unsafe { pool.deallocate(block.cast(), layout(1000)) }; } - }); - } - - #[test] - fn realloc_within_class_keeps_block() { - with_pool(|pool| { - let l = layout(100, 8); - let ptr = pool.alloc(l).unwrap(); - // SAFETY: realloc contract — `ptr` allocated with `l`. - let grown = unsafe { pool.realloc(ptr, l, 120) }.unwrap(); - assert_eq!(ptr, grown, "same class must realloc in place"); - // SAFETY: allocated above; 120 rounds to the same class as 100. - unsafe { pool.dealloc(grown, layout(120, 8)) }; - }); + } } #[test] - fn realloc_across_classes_preserves_contents() { - with_pool(|pool| { - let l = layout(64, 8); - let ptr = pool.alloc(l).unwrap(); - for i in 0..64u8 { - // SAFETY: live exclusive allocation of 64 bytes. - unsafe { ptr.as_ptr().add(usize::from(i)).write(i) }; - } - // SAFETY: realloc contract — `ptr` allocated with `l`. - let grown = unsafe { pool.realloc(ptr, l, 4096) }.unwrap(); - for i in 0..64u8 { - // SAFETY: `grown` is live for 4096 bytes. - let got = unsafe { grown.as_ptr().add(usize::from(i)).read() }; - assert_eq!(got, i); - } - // Shrink across classes, and from the large path back into a class. - // SAFETY: `grown` allocated with the 4096 layout above. - let shrunk = unsafe { pool.realloc(grown, layout(4096, 8), 8) }.unwrap(); - // SAFETY: `shrunk` is live for 8 bytes. - assert_eq!(unsafe { shrunk.as_ptr().read() }, 0); - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(shrunk, layout(8, 8)) }; - }); + fn custom_chunk_size_and_alignment_work() { + let pool = ChunkPool::::new_in(Global); + let block = pool.allocate(layout(10)).unwrap(); + assert_eq!(block.len(), 1024); + assert_eq!(block.cast::().as_ptr().addr() % 64, 0); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(10)) }; + + // Bigger than this pool's chunk size: passes through. + let big = pool.allocate(layout(2000)).unwrap(); + assert!(big.len() >= 2000); + // SAFETY: allocated above with the same layout. + unsafe { pool.deallocate(big.cast(), layout(2000)) }; } #[test] - fn alloc_zeroed_scrubs_recycled_blocks() { - with_pool(|pool| { - let l = layout(256, 8); - let dirty = pool.alloc(l).unwrap(); - // SAFETY: live exclusive allocation of 256 bytes. - unsafe { dirty.as_ptr().write_bytes(0xFF, l.size()) }; - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(dirty, l) }; - - let zeroed = pool.alloc_zeroed(l).unwrap(); - assert_eq!( - zeroed, dirty, - "must recycle the dirty block for this test to be meaningful" - ); - for i in 0..l.size() { - // SAFETY: live exclusive allocation of 256 bytes. - assert_eq!(unsafe { zeroed.as_ptr().add(i).read() }, 0); + fn concurrent_use_smoke() { + const OPS: usize = if cfg!(miri) { 40 } else { 500 }; + + let pool = pool(); + thread::scope(|scope| { + for t in 0..4_usize { + let pool = &pool; + scope.spawn(move || { + for i in 0..OPS { + let size = 1 + (i * 37 + t * 101) % (2 * CHUNK_SIZE); + let l = Layout::from_size_align(size, 8).unwrap(); + let block = pool.allocate(l).unwrap(); + assert!(block.len() >= size); + let marker = u8::try_from((t * 50 + i) % 251).unwrap(); + // SAFETY: exclusive fresh block of at least `size` + // bytes, allocated above with layout `l`. + unsafe { + block.cast::().as_ptr().write(marker); + block.cast::().as_ptr().add(size - 1).write(marker); + assert_eq!(block.cast::().as_ptr().read(), marker); + assert_eq!(block.cast::().as_ptr().add(size - 1).read(), marker); + pool.deallocate(block.cast(), l); + } + } + }); } - // SAFETY: allocated above with the same layout. - unsafe { pool.dealloc(zeroed, l) }; }); } - /// Cheap deterministic PRNG so the stress test needs no dependencies. - fn lcg(state: &mut u64) -> u64 { - *state = - state.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1_442_695_040_888_963_407); - *state >> 33 - } - - /// An owned allocation in flight between threads. - struct SendBlock(NonNull, Layout); - // SAFETY: SendBlock represents exclusive ownership of the block, which is - // transferred wholesale to the receiving thread. - unsafe impl Send for SendBlock {} - + /// The one test against the real kernel-backed default; everything else + /// runs against `Global` so Miri can check it. #[test] - fn concurrent_stress_with_cross_thread_frees() { - const THREADS: usize = if cfg!(miri) { 3 } else { 8 }; - const OPS: usize = if cfg!(miri) { 60 } else { 20_000 }; - - with_pool(|pool| { - thread::scope(|scope| { - // Each thread frees blocks allocated by its neighbor, so - // pushes and pops of the same list run on different threads. - // Bounded channels keep the number of in-flight blocks well - // under the pool's per-class capacity regardless of thread - // scheduling. - let (senders, receivers): (Vec<_>, Vec<_>) = - (0..THREADS).map(|_| mpsc::sync_channel::(64)).unzip(); - let mut senders_rotated: Vec<_> = senders.into_iter().map(Some).collect(); - senders_rotated.rotate_left(1); - - for (thread_idx, (receiver, sender)) in - receivers.into_iter().zip(&mut senders_rotated).enumerate() - { - let sender = sender.take().unwrap(); - let pool = &*pool; - scope.spawn(move || { - let mut rng = 0x9E37_79B9_7F4A_7C15_u64 ^ thread_idx as u64; - for _ in 0..OPS { - let size = 1 + usize::try_from(lcg(&mut rng)).unwrap() - % (3 * MAX_BLOCK_SIZE / 2); - let align = 1 << (lcg(&mut rng) % 7); - let l = layout(size, align); - let ptr = pool.alloc(l).unwrap(); - let marker = u8::try_from(lcg(&mut rng) & 0xFF).unwrap(); - // SAFETY: fresh exclusive allocation of `size` bytes. - unsafe { - ptr.as_ptr().write(marker); - ptr.as_ptr().add(size - 1).write(marker); - } - // SAFETY: we still exclusively own the block. - let (first, last) = - unsafe { (ptr.as_ptr().read(), ptr.as_ptr().add(size - 1).read()) }; - assert_eq!((first, last), (marker, marker), "block corrupted"); - if let Err(returned) = sender.try_send(SendBlock(ptr, l)) { - // Neighbor's queue is full (or it finished); - // free locally instead of blocking, which in - // a ring of senders could deadlock. - let (mpsc::TrySendError::Full(SendBlock(ptr, l)) - | mpsc::TrySendError::Disconnected(SendBlock(ptr, l))) = returned; - // SAFETY: allocated above with layout `l`. - unsafe { pool.dealloc(ptr, l) }; - } - // Drain what our own producer has sent so far, - // interleaving cross-thread pops with pushes. - while let Ok(SendBlock(ptr, l)) = receiver.try_recv() { - // SAFETY: the neighbor allocated `ptr` with - // `l` and transferred ownership over the - // channel. - unsafe { pool.dealloc(ptr, l) }; - } - } - drop(sender); - for SendBlock(ptr, l) in receiver { - // SAFETY: the neighbor allocated `ptr` with `l` - // and transferred ownership over the channel. - unsafe { pool.dealloc(ptr, l) }; - } - }); - } - }); - }); + #[cfg(not(miri))] + fn mmap_backed_pool_smoke() { + let pool = ChunkPool::::new(); + let block = pool.allocate(layout(100)).unwrap(); + assert_eq!(block.len(), CHUNK_SIZE); + let addr = block.cast::().as_ptr().addr(); + // SAFETY: fresh exclusive block of at least 100 bytes. + unsafe { block.cast::().as_ptr().write_bytes(0xCD, 100) }; + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(block.cast(), layout(100)) }; + let again = pool.allocate(layout(200)).unwrap(); + assert_eq!(again.cast::().as_ptr().addr(), addr); + // SAFETY: allocated above with a fitting layout. + unsafe { pool.deallocate(again.cast(), layout(200)) }; } } diff --git a/crates/fspy_alloc/src/slab.rs b/crates/fspy_alloc/src/slab.rs deleted file mode 100644 index 3aca476dc..000000000 --- a/crates/fspy_alloc/src/slab.rs +++ /dev/null @@ -1,203 +0,0 @@ -//! Slab memory layout and the [`Slab`] view type. -//! -//! Each size class carves fixed-size blocks out of [`SLAB_SIZE`]-byte, -//! `SLAB_SIZE`-aligned slabs. A slab looks like this: -//! -//! ```text -//! | SlabHeader | links: [AtomicU32; blocks] | pad to 4 KiB | block 0 | block 1 | ... | -//! ``` -//! -//! Free blocks are chained through the `links` side table — never through -//! block memory itself — so block payloads are only ever plain data. That -//! keeps every concurrent access well-defined under the memory model (no -//! mixed atomic/non-atomic reads of the same bytes) and Miri-clean. - -use core::{ - num::NonZero, - ptr::NonNull, - sync::atomic::{AtomicU32, AtomicU64}, -}; - -use crate::class::{CLASS_COUNT, MAX_POOL_ALIGN, block_size}; - -/// Slab size and alignment. Masking a block address with `!(SLAB_SIZE - 1)` -/// recovers its slab base. -pub const SLAB_SIZE: usize = 1 << 20; - -const fn round_up(value: usize, align: usize) -> usize { - (value + align - 1) & !(align - 1) -} - -/// Byte offset of the block area inside a slab holding `blocks` blocks. -const fn block_area_offset(blocks: usize) -> usize { - round_up(size_of::() + blocks * size_of::(), MAX_POOL_ALIGN) -} - -const fn compute_blocks_per_slab(class: usize) -> usize { - let size = block_size(class); - // Upper bound ignoring metadata, then shrink until header, link table, - // padding, and blocks all fit in one slab. - let mut blocks = SLAB_SIZE / size; - while block_area_offset(blocks) + blocks * size > SLAB_SIZE { - blocks -= 1; - } - blocks -} - -/// Blocks per slab, for each size class. -pub const BLOCKS_PER_SLAB: [usize; CLASS_COUNT] = { - let mut table = [0usize; CLASS_COUNT]; - let mut class = 0; - while class < CLASS_COUNT { - table[class] = compute_blocks_per_slab(class); - class += 1; - } - table -}; - -const _: () = { - assert!(size_of::() == 16); - assert!(MAX_POOL_ALIGN >= align_of::()); - let mut class = 0; - while class < CLASS_COUNT { - assert!(BLOCKS_PER_SLAB[class] > 0); - assert!( - block_area_offset(BLOCKS_PER_SLAB[class]) + BLOCKS_PER_SLAB[class] * block_size(class) - <= SLAB_SIZE - ); - class += 1; - } -}; - -/// Lives at the base of every slab mapping, ahead of the link table. -#[repr(C)] -struct SlabHeader { - /// This slab's index in its class's slab table. - slab_idx: u32, - _reserved: u32, - /// Number of blocks ever carved off this slab (monotonic; values at or - /// beyond the class's blocks-per-slab mean the slab is exhausted). 64-bit - /// so over-counting by racing threads can never wrap it around. - carved: AtomicU64, -} - -/// A view of one live slab mapping of a particular class. -/// -/// # Invariant -/// -/// `base` points to a [`SLAB_SIZE`]-byte, `SLAB_SIZE`-aligned mapping laid -/// out for `class` (initialized header, link table, block area) that is -/// never unmapped while the pool is in use. All raw-pointer arithmetic on -/// slab memory lives in this type's methods; the unsafe constructors are the -/// only places the invariant is asserted, and every accessor then relies on -/// it for bounds and liveness. -#[derive(Clone, Copy)] -pub struct Slab { - base: NonNull, - class: usize, -} - -impl Slab { - /// Wraps a slab pointer loaded from a class's slab table. - /// - /// # Safety - /// - /// `base` must be a non-null entry of the class's slab table. Such - /// entries are only published after [`Slab::init_header`] ran on a - /// suitable mapping, and are never cleared or unmapped while the pool is - /// in use, so the type invariant holds. - pub const unsafe fn from_published(base: NonNull, class: usize) -> Self { - Self { base, class } - } - - /// Recovers the slab containing a live pool block: slabs are - /// `SLAB_SIZE`-aligned, so masking the block address's low bits yields - /// the slab base (`with_addr` keeps the block pointer's provenance, which - /// covers the whole slab mapping it was carved from). Returns `None` only - /// for an address whose masked base would be null, which no real block - /// can produce. - /// - /// # Safety - /// - /// `block` must be a block of `class` previously handed out by this pool - /// and thus carved from a published slab of this class. - pub unsafe fn of_block(block: NonNull, class: usize) -> Option { - let base_addr = NonZero::new(block.as_ptr().addr() & !(SLAB_SIZE - 1))?; - Some(Self { base: block.with_addr(base_addr), class }) - } - - /// Initializes the header of a fresh slab mapping, making it publishable - /// into a slab table. - /// - /// # Safety - /// - /// `base` must point to a fresh, exclusive, zero-initialized, - /// `SLAB_SIZE`-byte and `SLAB_SIZE`-aligned mapping. The zero fill - /// doubles as the initial state of the carve counter and the link table. - pub unsafe fn init_header(base: NonNull, slab_idx: u32) { - // SAFETY: caller contract — the fresh exclusive mapping is aligned - // (SLAB_SIZE-aligned, far beyond SlabHeader's needs) and large - // enough for the header. - unsafe { - (*base.cast::().as_ptr()).slab_idx = slab_idx; - } - } - - const fn header(&self) -> &SlabHeader { - // SAFETY: type invariant — the mapping is live, SLAB_SIZE-aligned - // (far beyond SlabHeader's needs), and its header was initialized - // before publication; the only non-atomic field is never written - // again afterwards. - unsafe { self.base.cast::().as_ref() } - } - - /// This slab's index in its class's slab table. - pub const fn slab_idx(&self) -> usize { - self.header().slab_idx as usize - } - - /// The monotonic carve counter. - pub const fn carved(&self) -> &AtomicU64 { - &self.header().carved - } - - const fn blocks(&self) -> usize { - BLOCKS_PER_SLAB[self.class] - } - - /// The free-list link slot of `block_idx`. - pub fn link(&self, block_idx: usize) -> &AtomicU32 { - debug_assert!(block_idx < self.blocks()); - #[expect( - clippy::cast_ptr_alignment, - reason = "the link table starts at offset 16 of a SLAB_SIZE-aligned slab, so entries are 4-byte aligned" - )] - // SAFETY: type invariant plus the bound above put the slot within - // the slab's link table; slots are only ever accessed atomically, - // and the mapping outlives any borrow. - unsafe { - AtomicU32::from_ptr( - self.base.as_ptr().add(size_of::()).cast::().add(block_idx), - ) - } - } - - /// The address of block `block_idx`. - pub fn block(&self, block_idx: usize) -> NonNull { - debug_assert!(block_idx < self.blocks()); - let offset = block_area_offset(self.blocks()) + block_idx * block_size(self.class); - // SAFETY: type invariant plus the bound above keep `offset` within - // the SLAB_SIZE mapping. - unsafe { self.base.add(offset) } - } - - /// The index of a block previously returned by [`Slab::block`]. - pub fn block_index(&self, block: NonNull) -> usize { - let offset = - block.as_ptr().addr() - self.base.as_ptr().addr() - block_area_offset(self.blocks()); - debug_assert_eq!(offset % block_size(self.class), 0); - let block_idx = offset / block_size(self.class); - debug_assert!(block_idx < self.blocks()); - block_idx - } -} diff --git a/crates/fspy_alloc/src/sys.rs b/crates/fspy_alloc/src/sys.rs deleted file mode 100644 index 4fecdc085..000000000 --- a/crates/fspy_alloc/src/sys.rs +++ /dev/null @@ -1,25 +0,0 @@ -//! The backing-memory provider abstraction. -//! -//! The pool is generic over [`Sys`] so tests (and Miri) can substitute a -//! provider backed by the host allocator, while the real allocator obtains -//! anonymous mappings from the kernel (see the `mmap` module). - -use core::ptr::NonNull; - -/// Provides zero-initialized, aligned memory regions. -/// -/// Implementations must themselves be async-signal-safe and fork-safe: no -/// locks, no thread-locals, no libc malloc. -pub trait Sys { - /// Maps `size` bytes of zero-initialized memory aligned to `align` - /// (a power of two). Returns `None` when memory is exhausted. - fn map(size: usize, align: usize) -> Option>; - - /// Releases a region previously returned by [`Sys::map`]. - /// - /// # Safety - /// - /// `ptr` must have been returned by `Sys::map(size, align)` with the same - /// `size` and `align`, and must not be accessed afterwards. - unsafe fn unmap(ptr: NonNull, size: usize, align: usize); -} diff --git a/crates/fspy_alloc/tests/global_alloc.rs b/crates/fspy_alloc/tests/global_alloc.rs deleted file mode 100644 index ab0670ae5..000000000 --- a/crates/fspy_alloc/tests/global_alloc.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! Installs [`FspyAlloc`] as this test binary's global allocator, so every -//! allocation — including the test harness's own — exercises the allocator -//! end-to-end over real memory mappings. -#![cfg(all(unix, not(miri)))] // Miri covers the pool core via its mockable backend instead. - -use std::{collections::BTreeMap, sync::mpsc, thread}; - -use fspy_alloc::FspyAlloc; - -#[global_allocator] -static GLOBAL: FspyAlloc = FspyAlloc::new(); - -#[test] -fn collections_round_trip() { - let mut map = BTreeMap::new(); - for i in 0..1000_u32 { - let len = usize::try_from(i % 300).unwrap(); - map.insert(i, vec![0xAB_u8; len]); - } - assert_eq!(map.len(), 1000); - for (i, bytes) in &map { - assert_eq!(bytes.len(), usize::try_from(i % 300).unwrap()); - assert!(bytes.iter().all(|byte| *byte == 0xAB)); - } -} - -#[test] -fn vec_growth_reallocs_preserve_contents() { - let mut bytes = Vec::new(); - for i in 0..1_000_000_usize { - bytes.push(u8::try_from(i % 251).unwrap()); - } - for (i, byte) in bytes.iter().enumerate() { - assert_eq!(usize::from(*byte), i % 251); - } -} - -#[test] -fn large_and_zeroed_allocations() { - // Direct-mapped (beyond the largest size class), via the zeroing path. - let large = vec![0_u8; 5 * 1024 * 1024]; - assert!(large.iter().all(|byte| *byte == 0)); - // And through the plain path. - let boxed: Box<[u8]> = vec![7_u8; 300 * 1024].into_boxed_slice(); - assert!(boxed.iter().all(|byte| *byte == 7)); -} - -#[test] -fn many_small_allocations_span_slabs() { - // More live 16-byte-class blocks than one slab holds, forcing the pool - // through several slab installations (and the holding Vec through the - // direct-mapped path as it grows). - let boxes: Vec> = (0..120_000).map(|_| Box::new([0xEE_u8; 8])).collect(); - assert!(boxes.iter().all(|block| block.iter().all(|byte| *byte == 0xEE))); -} - -#[test] -fn threaded_producers_and_consumers() { - let threads = 8; - let per_thread = 5_000_usize; - let (sender, receiver) = mpsc::channel::>(); - thread::scope(|scope| { - for t in 0..threads { - let sender = sender.clone(); - scope.spawn(move || { - for i in 0..per_thread { - // Vary sizes across classes; blocks are freed (and often - // allocated) on the consumer thread below. - let len = 1 + (i * 37 + t * 101) % 5000; - sender.send(vec![u8::try_from(t % 251).unwrap(); len]).unwrap(); - } - }); - } - drop(sender); - let mut message_count = 0_usize; - for bytes in receiver { - message_count += 1; - assert!(!bytes.is_empty()); - let first = bytes[0]; - assert!(bytes.iter().all(|byte| *byte == first)); - } - assert_eq!(message_count, threads * per_thread); - }); -} diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 70a554563..7b54a06f5 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -8,6 +8,7 @@ publish = false crate-type = ["cdylib"] [target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true, features = ["alloc"] } anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index c42854b32..be1f2cc80 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "linux")] use std::ffi::CString; use std::{ - ffi::{CStr, OsStr}, + ffi::CStr, os::{fd::RawFd, unix::ffi::OsStrExt as _}, path::PathBuf, }; @@ -71,13 +71,27 @@ impl ToAbsolutePath for PathAt { if pathname.first().copied() == Some(b'/') { f(pathname.into()) } else { - let Some(mut abs_path) = get_fd_path(self.0)? else { + let Some(dir) = get_fd_path(self.0)? else { return f(None); }; + // Join `dir` and the relative `pathname` in a per-call bump + // arena instead of `PathBuf::push` on the global allocator. This + // runs on every fd-relative open/stat, including inside signal + // handlers and fork children, where the global allocator's locks + // are unsafe to take. The joined path only lives for the `f` + // call — nothing escapes the arena. + let arena = fspy_alloc::arena(); + let mut joined = allocator_api2::vec::Vec::new_in(&arena); + joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { - abs_path.push(OsStr::from_bytes(pathname)); + // Mirror `PathBuf::push`: exactly one separator between the + // (absolute, non-empty) dir and the relative pathname. + if joined.last() != Some(&b'/') { + joined.push(b'/'); + } + joined.extend_from_slice(pathname); } - f(Some(abs_path.as_os_str().as_bytes().as_bstr())) + f(Some(joined.as_bstr())) } } } diff --git a/crates/fspy_preload_unix/src/lib.rs b/crates/fspy_preload_unix/src/lib.rs index a9e667970..6c4e10b3e 100644 --- a/crates/fspy_preload_unix/src/lib.rs +++ b/crates/fspy_preload_unix/src/lib.rs @@ -1,16 +1,6 @@ // Compile as an empty crate on non-unix targets and on musl (where seccomp // alone handles access tracking). -/// This library interposes libc functions that POSIX declares -/// async-signal-safe (`open`, `stat`, `execve`, ...), so its own code may run -/// inside signal handlers and in the child of `fork()` in a multithreaded -/// process — contexts where taking the libc allocator's locks can deadlock. -/// Route every Rust allocation in this cdylib through fspy's lock-free, -/// mmap-backed allocator instead. -#[cfg(all(unix, not(target_env = "musl")))] -#[global_allocator] -static GLOBAL_ALLOCATOR: fspy_alloc::FspyAlloc = fspy_alloc::FspyAlloc::new(); - #[cfg(all(unix, not(target_env = "musl")))] mod client; #[cfg(all(unix, not(target_env = "musl")))] From 7f63a9316f62a28d7c4f73cbadf93d28eb00141f Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 11:06:01 +0800 Subject: [PATCH 3/4] feat(sigsafe): move the allocator into a crate of signal-safe syscall wrappers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The allocator is the first piece of a broader need: the preload library runs inside intercepted libc calls, so everything it uses must work in signal handlers, in the child of fork() in a multithreaded process, and before libc has finished initializing. sigsafe is where such code now lives. fspy_alloc becomes sigsafe::alloc, and the mmap/munmap/page-size calls it makes move behind sigsafe::mm and sigsafe::param, the first safe syscall wrappers the crate offers on its own. sigsafe promises that on Linux none of its calls go through libc, and enforces the promise at compile time: lib.rs references rustix::runtime, a module that exists only in rustix's raw-syscall backend, so any build configuration that selects rustix's libc backend — including a feature enabled outside this repository, which no build script can see — fails to compile instead of silently keeping libc in the picture. The README explains the purpose, the rules, and the enforcement. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 20 +-- Cargo.toml | 4 +- crates/fspy_alloc/Cargo.toml | 21 --- crates/fspy_preload_unix/Cargo.toml | 2 +- .../fspy_preload_unix/src/client/convert.rs | 2 +- crates/sigsafe/Cargo.toml | 35 +++++ crates/sigsafe/README.md | 60 +++++++ .../src => sigsafe/src/alloc}/mmap.rs | 19 +-- .../src/lib.rs => sigsafe/src/alloc/mod.rs} | 82 ++++++---- .../src => sigsafe/src/alloc}/pool.rs | 2 +- crates/sigsafe/src/lib.rs | 40 +++++ crates/sigsafe/src/mm.rs | 11 ++ crates/sigsafe/src/param.rs | 147 ++++++++++++++++++ 13 files changed, 365 insertions(+), 80 deletions(-) delete mode 100644 crates/fspy_alloc/Cargo.toml create mode 100644 crates/sigsafe/Cargo.toml create mode 100644 crates/sigsafe/README.md rename crates/{fspy_alloc/src => sigsafe/src/alloc}/mmap.rs (92%) rename crates/{fspy_alloc/src/lib.rs => sigsafe/src/alloc/mod.rs} (61%) rename crates/{fspy_alloc/src => sigsafe/src/alloc}/pool.rs (99%) create mode 100644 crates/sigsafe/src/lib.rs create mode 100644 crates/sigsafe/src/mm.rs create mode 100644 crates/sigsafe/src/param.rs diff --git a/Cargo.lock b/Cargo.lock index b783495c7..5b2ecd7dc 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1252,15 +1252,6 @@ dependencies = [ "winsafe 0.0.27", ] -[[package]] -name = "fspy_alloc" -version = "0.0.0" -dependencies = [ - "allocator-api2", - "bump-scope", - "rustix", -] - [[package]] name = "fspy_benchmark" version = "0.0.0" @@ -1318,11 +1309,11 @@ dependencies = [ "artifact_profile", "bstr", "ctor", - "fspy_alloc", "fspy_shared", "fspy_shared_unix", "libc", "nix 0.31.2", + "sigsafe", "wincode", ] @@ -3448,6 +3439,15 @@ dependencies = [ "libc", ] +[[package]] +name = "sigsafe" +version = "0.0.0" +dependencies = [ + "allocator-api2", + "bump-scope", + "rustix", +] + [[package]] name = "simd-adler32" version = "0.3.8" diff --git a/Cargo.toml b/Cargo.toml index f1c766099..457ba5cc3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -75,7 +75,6 @@ materialized_artifact = { path = "crates/materialized_artifact" } materialized_artifact_build = { path = "crates/materialized_artifact_build" } flate2 = "1.0.35" fspy = { path = "crates/fspy" } -fspy_alloc = { path = "crates/fspy_alloc" } fspy_benchmark_launcher = { path = "crates/fspy_benchmark_launcher", artifact = "bin" } fspy_benchmark_target = { path = "crates/fspy_benchmark_target", artifact = "bin" } fspy_detours_sys = { path = "crates/fspy_detours_sys" } @@ -123,7 +122,7 @@ ref-cast = "1.0.24" regex = "1.11.3" rusqlite = "0.39.0" rustc-hash = "2.1.1" -rustix = { version = "1", default-features = false, features = ["mm", "param", "use-libc-auxv"] } +rustix = { version = "1", default-features = false, features = ["mm"] } # SeccompAction::UserNotif (SECCOMP_RET_USER_NOTIF) was added after the latest published release (v0.5.0) seccompiler = { git = "https://github.com/rust-vmm/seccompiler", rev = "08587106340b8e3cb361c7561411510039436857" } serde = "1.0.219" @@ -131,6 +130,7 @@ serde_json = "1.0.140" serde_norway = "0.9.42" sha2 = "0.11.0" shell-escape = "0.1.5" +sigsafe = { path = "crates/sigsafe" } similar = "3.0.0" smallvec = { version = "2.0.0-alpha.12", features = ["std"] } snapshot_test = { path = "crates/snapshot_test" } diff --git a/crates/fspy_alloc/Cargo.toml b/crates/fspy_alloc/Cargo.toml deleted file mode 100644 index edcb89eaa..000000000 --- a/crates/fspy_alloc/Cargo.toml +++ /dev/null @@ -1,21 +0,0 @@ -[package] -name = "fspy_alloc" -edition = "2024" -license.workspace = true -publish = false - -[lib] -doctest = false - -[target.'cfg(unix)'.dependencies] -allocator-api2 = { workspace = true } -bump-scope = { workspace = true } -rustix = { workspace = true } - -[target.'cfg(unix)'.dev-dependencies] -# The `alloc` feature provides `Global`, letting tests run the pool against -# the host allocator (and thus under Miri). -allocator-api2 = { workspace = true, features = ["alloc"] } - -[lints] -workspace = true diff --git a/crates/fspy_preload_unix/Cargo.toml b/crates/fspy_preload_unix/Cargo.toml index 7b54a06f5..89ffaabdc 100644 --- a/crates/fspy_preload_unix/Cargo.toml +++ b/crates/fspy_preload_unix/Cargo.toml @@ -13,11 +13,11 @@ anyhow = { workspace = true } wincode = { workspace = true } bstr = { workspace = true, default-features = false } ctor = { workspace = true } -fspy_alloc = { workspace = true } fspy_shared = { workspace = true } fspy_shared_unix = { workspace = true } libc = { workspace = true } nix = { workspace = true, features = ["signal", "fs", "socket", "mman", "time"] } +sigsafe = { workspace = true } [build-dependencies] artifact_profile = { workspace = true } diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index be1f2cc80..70fc82a8d 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -80,7 +80,7 @@ impl ToAbsolutePath for PathAt { // handlers and fork children, where the global allocator's locks // are unsafe to take. The joined path only lives for the `f` // call — nothing escapes the arena. - let arena = fspy_alloc::arena(); + let arena = sigsafe::alloc::arena(); let mut joined = allocator_api2::vec::Vec::new_in(&arena); joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { diff --git a/crates/sigsafe/Cargo.toml b/crates/sigsafe/Cargo.toml new file mode 100644 index 000000000..9bd1af418 --- /dev/null +++ b/crates/sigsafe/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "sigsafe" +edition = "2024" +license.workspace = true +publish = false + +[lib] +doctest = false + +[target.'cfg(unix)'.dependencies] +allocator-api2 = { workspace = true } +bump-scope = { workspace = true } +rustix = { workspace = true } + +# On Linux the page size is probed from the kernel directly (see param.rs); +# rustix's `param` is only needed where sysconf is the platform interface. +[target.'cfg(all(unix, not(target_os = "linux")))'.dependencies] +rustix = { workspace = true, features = ["param"] } + +# The compile-time backend check in lib.rs needs a `linux_raw`-gated rustix +# item to reference; `runtime` is the module that has one. +[target.'cfg(target_os = "linux")'.dependencies] +rustix = { workspace = true, features = ["runtime"] } + +# Cross-validates the page-size probe against rustix's auxv-based answer. +[target.'cfg(target_os = "linux")'.dev-dependencies] +rustix = { workspace = true, features = ["param"] } + +[target.'cfg(unix)'.dev-dependencies] +# The `alloc` feature provides `Global`, letting tests run the pool against +# the host allocator (and thus under Miri). +allocator-api2 = { workspace = true, features = ["alloc"] } + +[lints] +workspace = true diff --git a/crates/sigsafe/README.md b/crates/sigsafe/README.md new file mode 100644 index 000000000..083c4adc6 --- /dev/null +++ b/crates/sigsafe/README.md @@ -0,0 +1,60 @@ +# sigsafe + +Unix syscall wrappers that are safe to call where libc is not. + +## Why this crate exists + +The fspy preload library injects itself into traced programs and intercepts +their libc calls (`open`, `stat`, `execve`, ...). POSIX declares those +functions async-signal-safe, so programs are allowed to call them: + +- inside a signal handler, +- in the child of `fork()` of a multithreaded program, +- while the process is still starting up, before libc is fully initialized. + +Most of libc is off limits in those places. `malloc` is the classic trap: a +signal can pause a thread while it holds malloc's lock, and `fork()` copies a +locked lock into a child that has no thread left to unlock it — the next +`malloc` waits forever. Interception code runs exactly there, so anything it +calls must work without libc's machinery, or the traced program can hang. + +## The rules + +Every function in this crate follows three rules: + +1. **Syscalls only.** On Linux, nothing goes through libc — the syscall + instructions are emitted directly (rustix's raw backend). On macOS there + is no stable syscall interface, so calls go through libSystem's wrappers; + for the calls exposed here those are thin stubs with no locks and no + state. +2. **No locks, no hidden state.** Nothing a signal or a `fork()` could catch + locked or half-written. Where shared state is unavoidable it is a fixed + set of atomics, each touched by single complete operations. +3. **No global allocation.** No function touches a heap behind the caller's + back. Code that needs memory gets it from an explicit allocator — + [`alloc`](src/alloc/mod.rs) provides one built on `mmap`. + +## How rule 1 is enforced on Linux + +rustix can be built with a libc backend instead of raw syscalls, and anything +in the dependency graph — including crates outside this repository — can +select it (the `rustix/use-libc` feature, or +`RUSTFLAGS=--cfg=rustix_use_libc`). No build script can detect that reliably, +so [`lib.rs`](src/lib.rs) checks at compile time instead: it references +`rustix::runtime`, a module that exists only in rustix's raw-syscall build. +Selecting the libc backend makes this crate fail to compile, rather than +silently losing the guarantee. + +## What's inside + +Functions whose rustix implementation already meets the rules are re-exposed +as-is; being listed in a module here is what marks a call as allowed, and the +backend check above is what keeps that true. + +- `mm` — anonymous memory mappings: `mmap_anonymous`, `munmap`. +- `param` — `page_size`. +- `alloc` — allocation without malloc: `alloc::arena()` gives one + intercepted call a bump arena that draws 64 KiB chunks from a process-wide + lock-free pool and returns them when the call ends. Taking or returning a + chunk is one atomic swap; when the pool is empty, chunks come straight + from the kernel through `mm`. diff --git a/crates/fspy_alloc/src/mmap.rs b/crates/sigsafe/src/alloc/mmap.rs similarity index 92% rename from crates/fspy_alloc/src/mmap.rs rename to crates/sigsafe/src/alloc/mmap.rs index 3cc09d876..40791d287 100644 --- a/crates/fspy_alloc/src/mmap.rs +++ b/crates/sigsafe/src/alloc/mmap.rs @@ -6,7 +6,11 @@ use core::{ }; use allocator_api2::alloc::{AllocError, Allocator}; -use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; + +use crate::{ + mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}, + param::page_size, +}; /// A stateless allocator: every allocation is a fresh anonymous mapping and /// every deallocation an `munmap`. @@ -16,10 +20,9 @@ use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; /// It holds no state at all — no locks, no free lists, no thread-locals; /// the kernel does all the bookkeeping. A signal or a `fork()` can never /// catch it holding a lock or a half-written structure, because there is -/// nothing to hold. Mapping syscalls go straight to the kernel on Linux -/// (rustix's raw backend) and through the thin libSystem wrappers on macOS. -/// The page size comes from static process data (`getauxval` on Linux, -/// `sysconf` on macOS) — no locks or allocation there either. +/// nothing to hold. The mapping calls and the page-size read come from +/// [`crate::mm`] and [`crate::param`], which carry the same guarantee (see +/// their docs). /// /// # What it accepts /// @@ -55,12 +58,6 @@ use rustix::mm::{MapFlags, ProtFlags, mmap_anonymous, munmap}; #[derive(Clone, Copy, Debug, Default)] pub struct MmapAllocator; -pub fn page_size() -> usize { - // getauxval on Linux, sysconf on macOS: reads of static process data — - // no locks, no allocation, valid from the first instruction on. - rustix::param::page_size() -} - // SAFETY: returned blocks are non-null, page-aligned (at least // `layout.align()` for every served layout), at least `layout.size()` bytes // large (the returned length reports the exact mapped size), stay valid diff --git a/crates/fspy_alloc/src/lib.rs b/crates/sigsafe/src/alloc/mod.rs similarity index 61% rename from crates/fspy_alloc/src/lib.rs rename to crates/sigsafe/src/alloc/mod.rs index 24b97020e..1a815f36b 100644 --- a/crates/fspy_alloc/src/lib.rs +++ b/crates/sigsafe/src/alloc/mod.rs @@ -1,24 +1,15 @@ -//! Async-signal-safe allocation for the fspy preload library. +//! Allocation that never touches libc malloc. //! -//! The preload library interposes libc functions that POSIX declares -//! async-signal-safe (`open`, `stat`, `execve`, ...). Its code therefore runs -//! inside signal handlers and in the child of `fork()` in a multithreaded -//! process — contexts where taking libc malloc's locks can deadlock on a -//! suspended or vanished lock holder, so the preload's own allocations must -//! never go through libc malloc. -//! -//! This crate stacks three layers, and exposes only the top one, [`arena`]. -//! `MmapAllocator` is the bottom: a stateless allocator where every -//! allocation is a fresh anonymous mapping from the kernel. `ChunkPool` -//! sits on top of it and caches fixed-size chunks, so that frequent short -//! tracing calls can reuse memory instead of paying two syscalls per call. -//! [`arena`] creates one `bump_scope::Bump` per intercepted call, drawing -//! its chunks from the process-wide pool and returning them on drop. - -// Compile as an empty crate on non-unix targets: the allocator backs the -// unix preload library. -#![cfg(unix)] -#![cfg_attr(not(test), no_std)] +//! Taking malloc's lock is the classic way for interposed code to deadlock a +//! traced program (see the crate docs), so the preload library allocates +//! through this module instead. It stacks three layers and exposes only the +//! top one, [`arena`]. `MmapAllocator` is the bottom: a stateless allocator +//! where every allocation is a fresh anonymous mapping from [`crate::mm`]. +//! `ChunkPool` sits on top of it and caches fixed-size chunks, so that +//! frequent short tracing calls can reuse memory instead of paying two +//! syscalls per call. [`arena`] creates one `bump_scope::Bump` per +//! intercepted call, drawing its chunks from the process-wide pool and +//! returning them on drop. mod mmap; mod pool; @@ -34,7 +25,8 @@ use pool::ChunkPool; /// Every cached chunk is 64 KiB: a whole multiple of the page size on all /// supported targets, and big enough that most intercepted calls fit their -/// allocations into a single chunk. +/// allocations into a single chunk. [`ArenaSettings`] pins the arenas' own +/// chunk sizing to this same value. const CHUNK_SIZE: usize = 64 * 1024; /// The alignment chunks are allocated with. Must be at least the alignment /// `bump_scope::Bump` uses for its chunk requests — 16 (see @@ -51,6 +43,22 @@ const SLOTS: usize = 64; /// first allocation on — even before any constructor has run. static CHUNK_POOL: ChunkPool = ChunkPool::new(); +/// The `Bump` settings the arenas use — the defaults, with two changes: +/// +/// - `WithGuaranteedAllocated`: an arena starts life without a chunk, +/// so creating one allocates nothing. +/// - `WithMinimumChunkSize`: the arena's first chunk request is +/// sized to the pool's chunks, making the coupling explicit — rather than +/// relying on the pool rounding the default 512-byte first request up to a +/// whole chunk. (Both end up serving the same memory: the pool answers any +/// request up to `CHUNK_SIZE` with a whole chunk, and `Bump` uses the full +/// returned length.) The request comes out slightly *under* `CHUNK_SIZE` — +/// bump-scope deducts an assumed allocator-header overhead from its +/// minimum — which is exactly what keeps a minimum-sized request within +/// the pool's `size <= CHUNK_SIZE` gate; the +/// `bump_chunk_requests_fit_the_pool_gates` test pins that fit. +type ArenaSettings = <::WithGuaranteedAllocated as BumpAllocatorSettings>::WithMinimumChunkSize; + /// `Bump::unallocated` requires its base allocator to implement `Default` /// (an arena without chunks has nowhere to store an allocator value, so it /// conjures one on first use). Point defaulted references at the @@ -77,12 +85,9 @@ impl Default for &'static ChunkPool impl Allocator { - // The default `Bump` settings, except the arena starts life without a - // chunk, so creating one allocates nothing. - type Settings = ::WithGuaranteedAllocated; Bump::< AllocatorApi2V02Compat<&'static ChunkPool>, - Settings, + ArenaSettings, >::unallocated() } @@ -95,17 +100,28 @@ mod tests { use super::*; /// Pins the fit between `Bump`'s chunk requests and the pool's gates - /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]). - /// bump-scope keeps those request parameters private, so this test is - /// the enforcement: it fails if an upgrade ever changes them. + /// (alignment at most [`CHUNK_ALIGN`], size at most [`CHUNK_SIZE`]), + /// under the same [`ArenaSettings`] the arenas use — including that a + /// request under `WithMinimumChunkSize` still fits the + /// `size <= CHUNK_SIZE` gate. bump-scope keeps its request parameters + /// private, so this test is the enforcement: it fails if an upgrade + /// ever changes them. + /// [`ArenaSettings`] with `GuaranteedAllocated` flipped back on: without + /// it, `Bump` demands a `Default` base allocator, and a reference to the + /// test's stack-local pool cannot provide one. Chunk request sizing — + /// what the test pins — is unaffected by that flag. + type TestSettings = ::WithGuaranteedAllocated; + #[test] fn bump_chunk_requests_fit_the_pool_gates() { let pool = ChunkPool::::new_in(Global); - // `try_new_in` requests the first chunk right away; if that request - // were over-aligned or oversized, the pool would refuse and this - // would be an error. - let bump: Bump>> = - Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); + // `try_new_in` requests the first chunk right away, at the settings' + // minimum chunk size; if that request were over-aligned or + // oversized, the pool would refuse and this would be an error. + let bump: Bump< + AllocatorApi2V02Compat<&ChunkPool>, + TestSettings, + > = Bump::try_new_in(AllocatorApi2V02Compat(&pool)).unwrap(); let block = bump.allocate(Layout::from_size_align(100, 8).unwrap()).unwrap(); let block_addr = block.cast::().as_ptr().addr(); drop(bump); diff --git a/crates/fspy_alloc/src/pool.rs b/crates/sigsafe/src/alloc/pool.rs similarity index 99% rename from crates/fspy_alloc/src/pool.rs rename to crates/sigsafe/src/alloc/pool.rs index f0e3e4112..3dcf31573 100644 --- a/crates/fspy_alloc/src/pool.rs +++ b/crates/sigsafe/src/alloc/pool.rs @@ -8,7 +8,7 @@ use core::{ use allocator_api2::alloc::{AllocError, Allocator}; -use crate::MmapAllocator; +use super::mmap::MmapAllocator; /// A fixed-size cache of memory chunks. /// diff --git a/crates/sigsafe/src/lib.rs b/crates/sigsafe/src/lib.rs new file mode 100644 index 000000000..9ee868345 --- /dev/null +++ b/crates/sigsafe/src/lib.rs @@ -0,0 +1,40 @@ +//! Unix syscall wrappers that are safe to call where libc is not: in signal +//! handlers, in the child of `fork()` in a multithreaded process, and before +//! libc has finished initializing. +//! +//! The fspy preload library interposes libc functions that POSIX declares +//! async-signal-safe (`open`, `stat`, `execve`, ...), so its code runs in all +//! of those places, where libc's own machinery — locks, lazy initialization, +//! malloc — is off limits. Everything this crate exposes follows three rules: +//! syscalls only (never through libc on Linux), no locks and no hidden state, +//! and no global allocation. See README.md for the full approach. + +// Compile as an empty crate on non-unix targets: the crate backs the unix +// preload library. +#![cfg(unix)] +#![cfg_attr(not(test), no_std)] + +pub mod alloc; +pub mod mm; +pub mod param; + +pub use rustix::io::Errno; + +// Compile-time proof that rustix uses its raw-syscall backend (`linux_raw`) +// on Linux — and with it, that no call in this crate goes through libc there. +// `rustix::runtime` is gated on that backend (`#[cfg(linux_raw)]`), so this +// reference fails to resolve, failing the whole build, whenever anything +// selects the libc backend instead: the `rustix/use-libc` feature (which any +// crate in the dependency graph can enable through feature unification, where +// no build script could ever see it), `RUSTFLAGS=--cfg=rustix_use_libc`, or a +// target rustix has no raw backend for. Miri also forces the libc backend, so +// it is exempted: it type-checks rather than ships code. +// +// The module's contents are not covered by rustix's stability promise, so a +// rustix upgrade may break this line. If that happens, re-point it at any +// other `rustix::runtime` item — do not delete it: it is the only enforcement +// of the no-libc rule above. +#[cfg(all(target_os = "linux", not(miri)))] +const _: () = { + let _ = rustix::runtime::exit_group; +}; diff --git a/crates/sigsafe/src/mm.rs b/crates/sigsafe/src/mm.rs new file mode 100644 index 000000000..c7cdecd05 --- /dev/null +++ b/crates/sigsafe/src/mm.rs @@ -0,0 +1,11 @@ +//! Anonymous memory mappings. +//! +//! Re-exposed from rustix as-is: these are single syscalls against the +//! kernel's own address-space bookkeeping — no libc state, no locks, no +//! allocation — so they already meet this crate's rules everywhere it +//! promises to work. What this module adds is the curation (being listed +//! here is what marks them safe for signal handlers, fork children, and +//! pre-libc startup) and the crate-level backend check, which guarantees +//! they cannot silently turn into libc calls on Linux. + +pub use rustix::mm::{MapFlags, MprotectFlags, ProtFlags, mmap_anonymous, mprotect, munmap}; diff --git a/crates/sigsafe/src/param.rs b/crates/sigsafe/src/param.rs new file mode 100644 index 000000000..0e907689a --- /dev/null +++ b/crates/sigsafe/src/param.rs @@ -0,0 +1,147 @@ +//! Process parameters. +//! +//! On Linux, `page_size` asks the kernel itself — this crate stands in for +//! libc, so it cannot lean on libc's startup knowledge. The kernel hands a +//! process its auxiliary vector exactly once, above the stack pointer at +//! entry, and only the code owning the entry point (a libc, or ld.so) sees +//! it; the retroactive interfaces (`prctl(PR_GET_AUXV)`, `/proc/self/auxv`) +//! need a 6.4 kernel or a mounted `/proc`. So instead of the auxiliary +//! vector, the value is learned from syscall behavior that any kernel +//! version guarantees: see [`linux::page_size`]. +//! +//! On other unix platforms (macOS, where every syscall goes through +//! libSystem by platform contract anyway) it is rustix's `sysconf`, a +//! lock-free read of startup data. + +#[cfg(target_os = "linux")] +pub use linux::page_size; +#[cfg(not(target_os = "linux"))] +pub use rustix::param::page_size; + +#[cfg(target_os = "linux")] +mod linux { + use core::{ + ptr, + sync::atomic::{AtomicUsize, Ordering}, + }; + + use crate::mm::{self, MapFlags, MprotectFlags, ProtFlags}; + + /// Smallest page size of any Linux configuration; the probe starts + /// here. + const MIN_PROBE_PAGE: usize = 4096; + /// Generous upper bound for the probe; no supported configuration uses + /// larger pages. + const MAX_PROBE_PAGE: usize = 1 << 20; + + /// Returns the size of a memory page, or 0 if it cannot be determined. + /// + /// Discovered from the kernel with a handful of raw syscalls, once per + /// process, and cached in an atomic: no libc, no auxiliary vector, no + /// `/proc`, no minimum kernel version, no allocation, and no init to + /// call first — safe from the first instruction on, in signal handlers, + /// and in the child of `fork()`. + /// + /// Zero — the kernel refusing a scratch mapping, i.e. out of memory — + /// is not cached, so a later call retries; callers treat it as "fail + /// this operation" (an allocator refuses the allocation). + #[must_use] + pub fn page_size() -> usize { + static CACHE: AtomicUsize = AtomicUsize::new(0); + let cached = CACHE.load(Ordering::Relaxed); + if cached != 0 { + return cached; + } + let page = probe_page_size(); + if page != 0 { + // Every probe returns the same value; the cache only avoids + // repeated probing. + CACHE.store(page, Ordering::Relaxed); + } + page + } + + /// Discovers the page size by probing, using nothing but raw syscalls. + /// + /// `mprotect` fails with `EINVAL` unless its address is a multiple of + /// the page size, and `mmap` returns page-aligned bases — both + /// kernel-guaranteed on every version. So re-protecting a scratch + /// mapping at increasing power-of-two offsets identifies the page size + /// as the first offset the kernel accepts: for a power-of-two page + /// size P, the first power of two P divides is P itself. + /// + /// Why not `rustix::param::page_size()` (fine on macOS, used there)? + /// Its Linux sources are exactly the ones this crate must not assume: + /// with `use-libc-auxv` it is libc's `sysconf`; without, its lazy init + /// needs `prctl(PR_GET_AUXV)` (kernel 6.4+) or `/proc/self/auxv`, it + /// panics when both are unavailable, and with rustix's `alloc` feature + /// enabled — which Cargo feature unification lets any other rustix + /// user in the build graph turn on for our copy — the `/proc` path + /// heap-allocates, which is forbidden in the contexts this crate + /// serves. The probe has none of those modes: no allocation, no panic, + /// no minimum kernel, and its worst case is 0. + #[cold] + fn probe_page_size() -> usize { + // Large enough that every probe below stays inside the mapping + // even after the kernel rounds the one-byte length up to a full + // page. + const SCRATCH_LEN: usize = 2 * MAX_PROBE_PAGE; + + // SAFETY: a fresh anonymous private mapping at no particular + // address has no memory-safety preconditions. + let Ok(base) = (unsafe { + mm::mmap_anonymous( + ptr::null_mut(), + SCRATCH_LEN, + ProtFlags::READ | ProtFlags::WRITE, + MapFlags::PRIVATE, + ) + }) else { + return 0; + }; + let mut page = 0; + let mut offset = MIN_PROBE_PAGE; + while offset <= MAX_PROBE_PAGE { + // SAFETY: `base + offset` (plus the page-rounded single byte) + // lies within the scratch mapping, which we exclusively own; + // the protection flags match the mapping's existing ones. + let accepted = unsafe { + mm::mprotect( + base.cast::().add(offset).cast(), + 1, + MprotectFlags::READ | MprotectFlags::WRITE, + ) + } + .is_ok(); + if accepted { + page = offset; + break; + } + offset *= 2; + } + // SAFETY: releasing the scratch mapping created above; nothing + // uses it afterwards. + let _ = unsafe { mm::munmap(base, SCRATCH_LEN) }; + page + } + + #[cfg(test)] + mod tests { + use super::*; + + /// Cross-validates the probe against rustix's auxv-based answer + /// (a dev-dependency here; CI Linux runners always have `/proc`). + #[test] + fn probe_matches_rustix() { + assert_eq!(page_size(), rustix::param::page_size()); + #[cfg(target_arch = "x86_64")] + assert_eq!(page_size(), 4096); + } + + #[test] + fn probe_is_stable_and_cached() { + assert_eq!(probe_page_size(), page_size()); + assert_eq!(page_size(), page_size()); + } + } +} From aa221e736d4ca16fafccc6dccc3b4f35183c2b20 Mon Sep 17 00:00:00 2001 From: wan9chi Date: Sun, 9 Aug 2026 14:25:35 +0800 Subject: [PATCH 4/4] feat(fspy): build exec's C pointer arrays in the per-call arena MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first use of the arena: `to_c_str_array` builds the NULL-terminated argv/envp pointer arrays that an intercepted exec hands to the real call, then drops them when that call returns — a temporary whose lifetime is already exactly a bump arena's. Building it must not go through libc malloc. Exec runs in the child of fork() in multithreaded programs (posix_spawn forks then execs), where malloc's lock may be held by a thread that no longer exists. Both platforms take this path, on every intercepted exec. The strings the array points at are still owned by `Exec` and still come from malloc; converting them, and the rest of the preload, is follow-up work. This change establishes the crate, the layers, and the lifetime discipline in the smallest place they all apply. Co-Authored-By: Claude Fable 5 --- .../fspy_preload_unix/src/client/convert.rs | 22 ++++--------------- .../fspy_preload_unix/src/client/raw_exec.rs | 11 ++++++++-- 2 files changed, 13 insertions(+), 20 deletions(-) diff --git a/crates/fspy_preload_unix/src/client/convert.rs b/crates/fspy_preload_unix/src/client/convert.rs index 70fc82a8d..c42854b32 100644 --- a/crates/fspy_preload_unix/src/client/convert.rs +++ b/crates/fspy_preload_unix/src/client/convert.rs @@ -1,7 +1,7 @@ #[cfg(target_os = "linux")] use std::ffi::CString; use std::{ - ffi::CStr, + ffi::{CStr, OsStr}, os::{fd::RawFd, unix::ffi::OsStrExt as _}, path::PathBuf, }; @@ -71,27 +71,13 @@ impl ToAbsolutePath for PathAt { if pathname.first().copied() == Some(b'/') { f(pathname.into()) } else { - let Some(dir) = get_fd_path(self.0)? else { + let Some(mut abs_path) = get_fd_path(self.0)? else { return f(None); }; - // Join `dir` and the relative `pathname` in a per-call bump - // arena instead of `PathBuf::push` on the global allocator. This - // runs on every fd-relative open/stat, including inside signal - // handlers and fork children, where the global allocator's locks - // are unsafe to take. The joined path only lives for the `f` - // call — nothing escapes the arena. - let arena = sigsafe::alloc::arena(); - let mut joined = allocator_api2::vec::Vec::new_in(&arena); - joined.extend_from_slice(dir.as_os_str().as_bytes()); if !pathname.is_empty() { - // Mirror `PathBuf::push`: exactly one separator between the - // (absolute, non-empty) dir and the relative pathname. - if joined.last() != Some(&b'/') { - joined.push(b'/'); - } - joined.extend_from_slice(pathname); + abs_path.push(OsStr::from_bytes(pathname)); } - f(Some(joined.as_bstr())) + f(Some(abs_path.as_os_str().as_bytes().as_bstr())) } } } diff --git a/crates/fspy_preload_unix/src/client/raw_exec.rs b/crates/fspy_preload_unix/src/client/raw_exec.rs index 250f9e282..a4d1a6336 100644 --- a/crates/fspy_preload_unix/src/client/raw_exec.rs +++ b/crates/fspy_preload_unix/src/client/raw_exec.rs @@ -43,10 +43,17 @@ impl RawExec { mut strs: Vec, f: impl FnOnce(*const *const libc::c_char) -> R, ) -> R { - let mut ptr_vec = Vec::<*const libc::c_char>::with_capacity(strs.len() + 1); + // The pointer array exists only for the `f` call below, and building + // it must not go through libc malloc: exec runs in the child of + // `fork()` in multithreaded programs (`posix_spawn` forks then + // execs), where malloc's lock may be held by a thread that no longer + // exists. A per-call arena has exactly this lifetime, and hands back + // the memory when the call ends. + let arena = sigsafe::alloc::arena(); + let mut ptr_vec = allocator_api2::vec::Vec::with_capacity_in(strs.len() + 1, &arena); for s in &mut strs { s.push(0); - ptr_vec.push(s.as_ptr().cast()); + ptr_vec.push(s.as_ptr().cast::()); } ptr_vec.push(null()); f(ptr_vec.as_ptr())