nordisk/znippy-zoomies

osm-katana/src/alloc_probe.rs

//! Test-only allocation counter — the instrument the tag-path guards read.
//!
//! Why a *thread-local* counter and not a global one: `cargo test` runs test
//! functions on several threads at once, so a process-wide `AtomicU64` would
//! count whatever the other tests happened to be doing and could never produce
//! a reproducible number. Each test thread counts only its own allocations, so
//! a measurement taken around a synchronous call is exact.
//!
//! Both cells are `const`-initialised and hold `Copy` types with no `Drop`, so
//! touching them from inside the allocator cannot itself allocate (a lazily
//! initialised or droppable thread-local would recurse into `alloc`).

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;

thread_local! {
    static ALLOCS: Cell<u64>  = const { Cell::new(0) };
    static ARMED:  Cell<bool> = const { Cell::new(false) };
}

pub struct CountingAlloc;

// SAFETY: every method forwards to `System` unchanged; the counter is a
// side-effect on a non-allocating thread-local.
unsafe impl GlobalAlloc for CountingAlloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        bump();
        unsafe { System.alloc(layout) }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { System.dealloc(ptr, layout) }
    }
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        bump();
        unsafe { System.realloc(ptr, layout, new_size) }
    }
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        bump();
        unsafe { System.alloc_zeroed(layout) }
    }
}

#[inline]
fn bump() {
    let _ = ARMED.try_with(|armed| {
        if armed.get() {
            let _ = ALLOCS.try_with(|n| n.set(n.get().saturating_add(1)));
        }
    });
}

#[global_allocator]
static GLOBAL: CountingAlloc = CountingAlloc;

/// Run `f` on this thread and return its value plus the number of heap
/// allocations (`alloc` + `alloc_zeroed` + `realloc`) it made **on this
/// thread**. Work handed to other threads is not counted, which is why every
/// guard using this drives a synchronous call.
pub fn count_allocs<T>(f: impl FnOnce() -> T) -> (T, u64) {
    ALLOCS.with(|n| n.set(0));
    ARMED.with(|a| a.set(true));
    let out = f();
    ARMED.with(|a| a.set(false));
    let n = ALLOCS.with(|n| n.get());
    (out, n)
}

#[cfg(test)]
mod self_tests {
    use super::count_allocs;

    /// The instrument must be able to read a non-zero — a counter that always
    /// says 0 would make every guard below it green forever.
    #[test]
    fn the_counter_counts() {
        let (_, quiet) = count_allocs(|| 1_u64 + 1);
        let (v, noisy) = count_allocs(|| {
            let mut v: Vec<u64> = Vec::new();
            for i in 0..1_000 {
                v.push(i);
            }
            v
        });
        assert_eq!(v.len(), 1_000);
        assert_eq!(quiet, 0, "arithmetic must not allocate");
        assert!(
            noisy >= 8,
            "1000 pushes must show growth reallocs, saw {noisy}"
        );
    }
}