loam_time/
alloc.rs

1use std::alloc::{GlobalAlloc, Layout};
2use std::cell::Cell;
3use std::sync::atomic::{AtomicBool, Ordering};
4use std::thread::LocalKey;
5
6thread_local! {
7    static ALLOC_BYTES: Cell<u64> = const { Cell::new(0) };
8    static DEALLOC_BYTES: Cell<u64> = const { Cell::new(0) };
9    static ALLOC_COUNT: Cell<u64> = const { Cell::new(0) };
10    static DEALLOC_COUNT: Cell<u64> = const { Cell::new(0) };
11}
12
13static ALLOC_INSTALLED: AtomicBool = AtomicBool::new(false);
14
15fn bump(counter: &'static LocalKey<Cell<u64>>, by: u64) {
16    let _ = counter.try_with(|count| count.set(count.get().wrapping_add(by)));
17}
18
19/// Counts requested bytes on the calling thread; allocator overhead is excluded.
20pub struct CountingAllocator<A: GlobalAlloc> {
21    inner: A,
22}
23
24impl<A: GlobalAlloc> CountingAllocator<A> {
25    pub const fn new(inner: A) -> Self {
26        Self { inner }
27    }
28}
29
30// SAFETY: Methods preserve the inner allocator contract; const TLS and wrapping Cell updates cannot unwind.
31unsafe impl<A: GlobalAlloc> GlobalAlloc for CountingAllocator<A> {
32    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
33        ALLOC_INSTALLED.store(true, Ordering::Relaxed);
34        bump(&ALLOC_COUNT, 1);
35        bump(&ALLOC_BYTES, layout.size() as u64);
36        // SAFETY: The caller supplies a valid nonzero allocation layout.
37        unsafe { self.inner.alloc(layout) }
38    }
39
40    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
41        bump(&DEALLOC_COUNT, 1);
42        bump(&DEALLOC_BYTES, layout.size() as u64);
43        // SAFETY: The caller supplies a live inner allocation and its original layout.
44        unsafe { self.inner.dealloc(ptr, layout) };
45    }
46
47    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
48        ALLOC_INSTALLED.store(true, Ordering::Relaxed);
49        bump(&ALLOC_COUNT, 1);
50        bump(&ALLOC_BYTES, layout.size() as u64);
51        // SAFETY: The caller supplies a valid nonzero allocation layout.
52        unsafe { self.inner.alloc_zeroed(layout) }
53    }
54
55    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
56        ALLOC_INSTALLED.store(true, Ordering::Relaxed);
57        bump(&DEALLOC_COUNT, 1);
58        bump(&DEALLOC_BYTES, layout.size() as u64);
59        bump(&ALLOC_COUNT, 1);
60        bump(&ALLOC_BYTES, new_size as u64);
61        // SAFETY: The caller supplies a live inner allocation, its layout, and a valid new size.
62        unsafe { self.inner.realloc(ptr, layout, new_size) }
63    }
64}
65
66#[derive(Clone, Copy, Debug, Default)]
67pub struct AllocSnapshot {
68    pub alloc_bytes: u64,
69    pub dealloc_bytes: u64,
70    pub alloc_count: u64,
71    pub dealloc_count: u64,
72}
73
74#[derive(Clone, Copy, Debug, Default)]
75pub struct AllocDelta {
76    pub net_bytes: i64,
77    pub alloc_bytes: u64,
78    pub alloc_count: u64,
79    pub dealloc_count: u64,
80}
81
82/// The calling thread's counters; `None` precedes the first allocation anywhere in the process.
83pub fn current_snapshot() -> Option<AllocSnapshot> {
84    if !ALLOC_INSTALLED.load(Ordering::Relaxed) {
85        return None;
86    }
87    Some(AllocSnapshot {
88        alloc_bytes: ALLOC_BYTES.with(Cell::get),
89        dealloc_bytes: DEALLOC_BYTES.with(Cell::get),
90        alloc_count: ALLOC_COUNT.with(Cell::get),
91        dealloc_count: DEALLOC_COUNT.with(Cell::get),
92    })
93}
94
95pub fn bytes_allocated_by(body: impl FnOnce()) -> Option<u64> {
96    if !ALLOC_INSTALLED.load(Ordering::Relaxed) {
97        body();
98        return None;
99    }
100    let before = ALLOC_BYTES.with(Cell::get);
101    body();
102    Some(ALLOC_BYTES.with(Cell::get).wrapping_sub(before))
103}
104
105/// `start` must be the earlier snapshot.
106pub fn delta(start: AllocSnapshot, end: AllocSnapshot) -> AllocDelta {
107    let alloc_bytes = end.alloc_bytes.saturating_sub(start.alloc_bytes);
108    let dealloc_bytes = end.dealloc_bytes.saturating_sub(start.dealloc_bytes);
109    AllocDelta {
110        net_bytes: (alloc_bytes as i64).saturating_sub(dealloc_bytes as i64),
111        alloc_bytes,
112        alloc_count: end.alloc_count.saturating_sub(start.alloc_count),
113        dealloc_count: end.dealloc_count.saturating_sub(start.dealloc_count),
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn delta_handles_dealloc_dominant() {
123        let start = AllocSnapshot {
124            alloc_bytes: 1_000,
125            dealloc_bytes: 200,
126            alloc_count: 10,
127            dealloc_count: 3,
128        };
129        let end = AllocSnapshot {
130            alloc_bytes: 1_100,
131            dealloc_bytes: 1_000,
132            alloc_count: 12,
133            dealloc_count: 20,
134        };
135        let d = delta(start, end);
136        assert_eq!(d.alloc_bytes, 100);
137        assert_eq!(d.net_bytes, -700);
138    }
139}
loam_time/alloc.rs

loam_time/
alloc.rs
