densa_alloc/
lib.rs

1//! Per-thread allocation counting, to *prove* hot paths never touch the heap.
2//!
3//! Install [`CountingAllocator`] as the global allocator of a binary or test crate:
4//!
5//! ```ignore
6//! #[global_allocator]
7//! static ALLOCATOR: densa_alloc::CountingAllocator = densa_alloc::CountingAllocator::system();
8//! ```
9//!
10//! then wrap the code under test in [`measure`]. Counters are per thread, so other threads
11//! (test harness, workers) never pollute a measurement. Without the allocator installed all
12//! counters stay at zero; always check [`is_active`] first so a zero-allocation assertion
13//! cannot pass vacuously.
14
15use std::alloc::{GlobalAlloc, Layout, System};
16use std::cell::Cell;
17use std::hint::black_box;
18
19/// Heap operations observed on one thread.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub struct AllocStats {
22    /// `alloc` + `alloc_zeroed` calls.
23    pub allocs: u64,
24    /// `dealloc` calls.
25    pub deallocs: u64,
26    /// `realloc` calls.
27    pub reallocs: u64,
28    /// Bytes requested by `alloc`, `alloc_zeroed` and `realloc` (new size).
29    pub bytes_allocated: u64,
30}
31
32impl AllocStats {
33    /// No operations: the starting point for counters and deltas.
34    pub const ZERO: Self = Self {
35        allocs: 0,
36        deallocs: 0,
37        reallocs: 0,
38        bytes_allocated: 0,
39    };
40
41    /// `true` if no heap operation of any kind happened.
42    pub fn is_zero(&self) -> bool {
43        self.allocs == 0 && self.deallocs == 0 && self.reallocs == 0
44    }
45
46    /// Operations that happened between `earlier` and `self`.
47    pub fn since(self, earlier: Self) -> Self {
48        Self {
49            allocs: self.allocs.wrapping_sub(earlier.allocs),
50            deallocs: self.deallocs.wrapping_sub(earlier.deallocs),
51            reallocs: self.reallocs.wrapping_sub(earlier.reallocs),
52            bytes_allocated: self.bytes_allocated.wrapping_sub(earlier.bytes_allocated),
53        }
54    }
55}
56
57thread_local! {
58    // `const` init + a type without `Drop`: lives in static TLS, so touching it from inside
59    // the allocator never allocates and never recurses into the allocator.
60    static STATS: Cell<AllocStats> = const { Cell::new(AllocStats::ZERO) };
61}
62
63#[inline]
64fn record(update: impl FnOnce(&mut AllocStats)) {
65    // `try_with` instead of `with`: never panic inside the allocator, even during
66    // thread teardown.
67    let _ = STATS.try_with(|cell| {
68        let mut stats = cell.get();
69        update(&mut stats);
70        cell.set(stats);
71    });
72}
73
74/// Snapshot of the calling thread's cumulative counters.
75pub fn thread_stats() -> AllocStats {
76    STATS.try_with(Cell::get).unwrap_or(AllocStats::ZERO)
77}
78
79/// Runs `f` and returns its result with the heap operations it performed on this thread.
80pub fn measure<R>(f: impl FnOnce() -> R) -> (R, AllocStats) {
81    let before = thread_stats();
82    let result = f();
83    (result, thread_stats().since(before))
84}
85
86/// `true` if [`CountingAllocator`] is the active global allocator.
87///
88/// Performs one tiny allocation and checks that it was counted.
89pub fn is_active() -> bool {
90    let ((), stats) = measure(|| drop(black_box(Box::new(0u8))));
91    stats.allocs > 0
92}
93
94/// A [`GlobalAlloc`] wrapper that counts every heap operation per thread.
95#[derive(Debug, Default)]
96pub struct CountingAllocator<A = System> {
97    inner: A,
98}
99
100impl CountingAllocator<System> {
101    /// Counting wrapper around the system allocator.
102    pub const fn system() -> Self {
103        Self { inner: System }
104    }
105}
106
107impl<A> CountingAllocator<A> {
108    /// Counting wrapper around any other allocator.
109    pub const fn new(inner: A) -> Self {
110        Self { inner }
111    }
112}
113
114#[inline]
115fn bytes(size: usize) -> u64 {
116    // Lossless on every supported target (usize <= 64 bits).
117    size as u64
118}
119
120// SAFETY: every method forwards its arguments unchanged to `inner`, so `inner`'s
121// `GlobalAlloc` guarantees carry over. The bookkeeping only updates a const-initialised,
122// `Drop`-free thread-local `Cell`, which neither allocates nor re-enters the allocator.
123unsafe impl<A: GlobalAlloc> GlobalAlloc for CountingAllocator<A> {
124    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
125        record(|s| {
126            s.allocs += 1;
127            s.bytes_allocated += bytes(layout.size());
128        });
129        // SAFETY: the caller upholds `GlobalAlloc::alloc`'s contract for `layout`.
130        unsafe { self.inner.alloc(layout) }
131    }
132
133    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
134        record(|s| {
135            s.allocs += 1;
136            s.bytes_allocated += bytes(layout.size());
137        });
138        // SAFETY: the caller upholds `GlobalAlloc::alloc_zeroed`'s contract for `layout`.
139        unsafe { self.inner.alloc_zeroed(layout) }
140    }
141
142    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
143        record(|s| s.deallocs += 1);
144        // SAFETY: the caller guarantees `ptr` was allocated by this allocator (hence by
145        // `inner`, since we forward every allocation) with `layout`.
146        unsafe { self.inner.dealloc(ptr, layout) }
147    }
148
149    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
150        record(|s| {
151            s.reallocs += 1;
152            s.bytes_allocated += bytes(new_size);
153        });
154        // SAFETY: the caller upholds `GlobalAlloc::realloc`'s contract; `ptr` came from
155        // `inner` because every allocation is forwarded to it.
156        unsafe { self.inner.realloc(ptr, layout, new_size) }
157    }
158}
densa_alloc/lib.rs

densa_alloc/
lib.rs
