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..5532d8734 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", "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 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_benchmark/src/main.rs b/crates/fspy_benchmark/src/main.rs index 91a899e1e..17f69e6a4 100644 --- a/crates/fspy_benchmark/src/main.rs +++ b/crates/fspy_benchmark/src/main.rs @@ -51,6 +51,10 @@ struct Suite { /// Unmeasured iterations run first, to fill caches and settle the runner. warmup: usize, metric: Metric, + /// Whether the target opens a relative path (the launcher's + /// `--relative`), driving the tracker's working-directory resolution and + /// path joining instead of the borrow-only absolute lane. + relative: bool, } /// Opens nothing, so the whole launch is the cost of starting a tracked @@ -62,12 +66,30 @@ const LAUNCH_SUITE: Suite = Suite { iterations: if cfg!(windows) { 150 } else { 300 }, warmup: 5, metric: Metric::Wall, + relative: false, }; /// Opens timed from inside the target, so they price interception rather than -/// the launch around it. -const ACCESS_SUITE: Suite = - Suite { name: "access", opens: "2048", iterations: 102, warmup: 3, metric: Metric::Typical }; +/// the launch around it. The absolute path takes the tracker's borrow-only +/// lane; the relative variant prices working-directory resolution and path +/// joining on top. +const ACCESS_SUITE: Suite = Suite { + name: "access", + opens: "2048", + iterations: 102, + warmup: 3, + metric: Metric::Typical, + relative: false, +}; + +const RELATIVE_ACCESS_SUITE: Suite = Suite { + name: "access-relative", + opens: "2048", + iterations: 102, + warmup: 3, + metric: Metric::Typical, + relative: true, +}; struct Backend { name: &'static str, @@ -86,11 +108,13 @@ fn main() { let backends = [Backend { name: "dynamic", target: DYNAMIC_TARGET }]; for backend in &backends { - validate(HEAD_LAUNCHER.as_ref(), backend.target); - if let Some(base_launcher) = &base_launcher { - validate(base_launcher, backend.target); + for relative in [false, true] { + validate(HEAD_LAUNCHER.as_ref(), backend.target, relative); + if let Some(base_launcher) = &base_launcher { + validate(base_launcher, backend.target, relative); + } } - for suite in [&LAUNCH_SUITE, &ACCESS_SUITE] { + for suite in [&LAUNCH_SUITE, &ACCESS_SUITE, &RELATIVE_ACCESS_SUITE] { run_suite(backend, suite, base_launcher.as_deref()); } } @@ -209,6 +233,9 @@ fn launch(launcher: &OsStr, mode: Option<&str>, backend: &Backend, suite: &Suite if let Some(mode) = mode { command.arg(mode); } + if suite.relative { + command.arg("--relative"); + } let output = command .args([backend.target, THREADS, suite.opens]) .stdin(Stdio::null()) @@ -228,9 +255,13 @@ fn launch(launcher: &OsStr, mode: Option<&str>, backend: &Backend, suite: &Suite } } -fn validate(launcher: &OsStr, target: &str) { - let status = Command::new(launcher) - .arg("--validate") +fn validate(launcher: &OsStr, target: &str, relative: bool) { + let mut command = Command::new(launcher); + command.arg("--validate"); + if relative { + command.arg("--relative"); + } + let status = command .arg(target) // One thread opening one batch: the count matches the target's // OPENS_PER_SAMPLE. Fewer would open nothing, which validation diff --git a/crates/fspy_benchmark_launcher/src/main.rs b/crates/fspy_benchmark_launcher/src/main.rs index 8a0a6c70e..80a54bb90 100644 --- a/crates/fspy_benchmark_launcher/src/main.rs +++ b/crates/fspy_benchmark_launcher/src/main.rs @@ -22,28 +22,42 @@ const MISSING_PATH: &str = "/.fspy-benchmark-missing"; #[cfg(windows)] const MISSING_PATH: &str = r"C:\.fspy-benchmark-missing"; +/// With `--relative`, the target opens this bare name from the filesystem +/// root instead, driving the tracker's relative-path lane: it must resolve +/// the working directory and join the two. Root as the working directory +/// makes the joined result exactly [`MISSING_PATH`], so validation checks +/// the same captured path in both modes. +const MISSING_RELATIVE_PATH: &str = ".fspy-benchmark-missing"; +#[cfg(unix)] +const ROOT_DIR: &str = "/"; +#[cfg(windows)] +const ROOT_DIR: &str = r"C:\"; + fn main() { let mut args = env::args_os().skip(1).collect::>(); - let mode = match args.first().map(OsString::as_os_str) { - Some(arg) if arg == "--untracked" => { - args.remove(0); - Mode::Untracked - } - Some(arg) if arg == "--validate" => { - args.remove(0); - Mode::Validate + let mut mode = Mode::Tracked; + let mut relative = false; + while let Some(flag) = args.first().map(OsString::as_os_str) { + if flag == "--untracked" { + mode = Mode::Untracked; + } else if flag == "--validate" { + mode = Mode::Validate; + } else if flag == "--relative" { + relative = true; + } else { + break; } - _ => Mode::Tracked, - }; + args.remove(0); + } let (target, target_args) = - args.split_first().expect("usage: fspy_benchmark_launcher [MODE] TARGET ARGS..."); + args.split_first().expect("usage: fspy_benchmark_launcher [FLAGS] TARGET ARGS..."); let runtime = Builder::new_multi_thread().worker_threads(2).enable_all().build().unwrap(); runtime.block_on(async { match mode { - Mode::Tracked => report(run_tracked(target, target_args).await).await, - Mode::Untracked => report(run_untracked(target, target_args).await).await, - Mode::Validate => validate(target, target_args).await, + Mode::Tracked => report(run_tracked(target, target_args, relative).await).await, + Mode::Untracked => report(run_untracked(target, target_args, relative).await).await, + Mode::Validate => validate(target, target_args, relative).await, } }); } @@ -62,14 +76,17 @@ struct Launch { /// Times the launch from just before the spawn to just after the wait, so /// that a tracked launch covers session setup, injection, and teardown, and /// nothing of this launcher's own startup. -async fn run_tracked(target: &OsString, target_args: &[OsString]) -> Launch { +async fn run_tracked(target: &OsString, target_args: &[OsString], relative: bool) -> Launch { let mut command = Command::new(target); command .args(target_args) - .arg(MISSING_PATH) + .arg(if relative { MISSING_RELATIVE_PATH } else { MISSING_PATH }) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()); + if relative { + command.current_dir(ROOT_DIR); + } let start = Instant::now(); let mut child = command .spawn(CancellationToken::new()) @@ -86,14 +103,17 @@ async fn run_tracked(target: &OsString, target_args: &[OsString]) -> Launch { Launch { wall_nanos, stdout } } -async fn run_untracked(target: &OsString, target_args: &[OsString]) -> Launch { +async fn run_untracked(target: &OsString, target_args: &[OsString], relative: bool) -> Launch { let mut command = tokio::process::Command::new(target); command .args(target_args) - .arg(MISSING_PATH) + .arg(if relative { MISSING_RELATIVE_PATH } else { MISSING_PATH }) .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()); + if relative { + command.current_dir(ROOT_DIR); + } let start = Instant::now(); let mut child = command.spawn().expect("failed to spawn untracked benchmark target"); let stdout = child.stdout.take().expect("untracked benchmark target has no stdout"); @@ -117,14 +137,20 @@ async fn report(mut launch: Launch) { /// Runs the target tracked and asserts that its accesses were captured, so /// that the harness never benchmarks tracking that silently stopped working. -async fn validate(target: &OsString, target_args: &[OsString]) { +/// In relative mode the captured path must come out identical — the tracker +/// resolves the root working directory and joins the bare name back into +/// [`MISSING_PATH`] — so the assertion below covers both modes. +async fn validate(target: &OsString, target_args: &[OsString], relative: bool) { let mut command = Command::new(target); command .args(target_args) - .arg(MISSING_PATH) + .arg(if relative { MISSING_RELATIVE_PATH } else { MISSING_PATH }) .stdin(Stdio::null()) .stdout(Stdio::null()) .stderr(Stdio::inherit()); + if relative { + command.current_dir(ROOT_DIR); + } let termination = command .spawn(CancellationToken::new()) .await 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")))]