concinnity_core/memory/
tracking.rs

1// The `GlobalAlloc` wrapper that counts what the Rust heap is holding.
2//
3// Installing it is the binary's call, since `#[global_allocator]` is a
4// per-program item and an embedder may want an allocator of its own:
5//
6//     #[global_allocator]
7//     static ALLOC: TrackingAlloc<System> = TrackingAlloc::new(System);
8//
9// Until some binary does, `stats()` reports `None` rather than a zero that would
10// read as "the heap is empty".
11
12use core::alloc::{GlobalAlloc, Layout};
13
14use crate::memory::counters::Counters;
15
16pub(crate) static COUNTERS: Counters = Counters::new();
17
18/// A `GlobalAlloc` that counts allocations and forwards them to `A`.
19pub struct TrackingAlloc<A> {
20    inner: A,
21    // The block this wrapper counts into, which for an installed allocator is
22    // always the process-global one. Naming it rather than reaching for the
23    // static is what lets a test weigh a wrapper of its own on its own scale.
24    counters: &'static Counters,
25}
26
27impl<A> TrackingAlloc<A> {
28    /// Wrap `inner` so every allocation through it is counted.
29    pub const fn new(inner: A) -> Self {
30        Self {
31            inner,
32            counters: &COUNTERS,
33        }
34    }
35
36    // A wrapper counting into `counters` instead of the process-global block.
37    #[cfg(test)]
38    pub(crate) const fn with_counters(inner: A, counters: &'static Counters) -> Self {
39        Self { inner, counters }
40    }
41}
42
43// SAFETY: every method forwards to `inner`, which upholds the `GlobalAlloc`
44// contract, and returns its pointer unchanged. The counter updates are relaxed
45// atomics over statically allocated storage, so the allocator cannot re-enter
46// itself.
47unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAlloc<A> {
48    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
49        // SAFETY: `layout` is forwarded unchanged from our caller, who upholds
50        // the `GlobalAlloc::alloc` contract.
51        let ptr = unsafe { self.inner.alloc(layout) };
52        if !ptr.is_null() {
53            self.counters.record_alloc(layout.size());
54            crate::memory::detail::record_alloc(layout.size());
55        }
56        ptr
57    }
58
59    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
60        // SAFETY: as `alloc`, forwarding our caller's obligations to `inner`.
61        let ptr = unsafe { self.inner.alloc_zeroed(layout) };
62        if !ptr.is_null() {
63            self.counters.record_alloc(layout.size());
64            crate::memory::detail::record_alloc(layout.size());
65        }
66        ptr
67    }
68
69    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
70        // SAFETY: `ptr` and `layout` come from our caller, who guarantees the
71        // block was allocated by us with this layout; we allocated it through
72        // `inner`, so `inner` is the right allocator to free it.
73        unsafe { self.inner.dealloc(ptr, layout) };
74        self.counters.record_free(layout.size());
75        crate::memory::detail::record_free(layout.size());
76    }
77
78    // Forwarded rather than left to the default alloc-copy-dealloc so `inner`
79    // can grow a block in place. A failed resize leaves the original block
80    // allocated and unchanged, so the counters only move on success.
81    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
82        // SAFETY: `ptr`, `layout` and `new_size` are forwarded unchanged from
83        // our caller, who upholds the `GlobalAlloc::realloc` contract.
84        let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) };
85        if !new_ptr.is_null() {
86            self.counters.record_realloc(layout.size(), new_size);
87            crate::memory::detail::record_realloc(layout.size(), new_size);
88        }
89        new_ptr
90    }
91}
92
93/// Install [`TrackingAlloc`] over the system allocator as this program's
94/// `#[global_allocator]`.
95///
96/// `#[global_allocator]` is a per-program item, so every binary, test binary,
97/// and benchmark that should count its own heap invokes this once at its crate
98/// root. A program without it runs on Rust's default allocator and reports no
99/// memory at all, which is what [`crate::memory::stats`] returning `None` means.
100///
101/// ```
102/// concinnity_core::install_global_allocator!();
103/// ```
104#[macro_export]
105macro_rules! install_global_allocator {
106    () => {
107        #[global_allocator]
108        static CN_GLOBAL_ALLOC: $crate::memory::TrackingAlloc<std::alloc::System> =
109            $crate::memory::TrackingAlloc::new(std::alloc::System);
110    };
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use core::ptr;
117    use std::alloc::System;
118
119    const SIZE: usize = 4096;
120    const GROWN: usize = 64 * 1024;
121    const SHRUNK: usize = 1024;
122
123    fn layout(size: usize) -> Layout {
124        Layout::from_size_align(size, 8).expect("valid layout")
125    }
126
127    // An inner allocator that always fails, for the paths that must count
128    // nothing.
129    struct Failing;
130
131    // SAFETY: `alloc` hands out no memory, returning the null `GlobalAlloc`
132    // defines as failure, so `dealloc` has no block it can be called with.
133    unsafe impl GlobalAlloc for Failing {
134        unsafe fn alloc(&self, _layout: Layout) -> *mut u8 {
135            ptr::null_mut()
136        }
137
138        unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
139            unreachable!("this allocator never hands out a block to free")
140        }
141    }
142
143    // Drives the real wrapper against the real system allocator, counting into
144    // a block of this test's own. The process-global counters are a net figure
145    // that a free on any other test thread moves between two reads, so they
146    // support only a threshold; a block nothing else counts into is exact.
147    #[test]
148    fn tracking_alloc_forwards_and_counts() {
149        static BLOCK: Counters = Counters::new();
150        let alloc = TrackingAlloc::with_counters(System, &BLOCK);
151        let layout = layout(SIZE);
152        assert_eq!(BLOCK.snapshot(), None, "nothing has allocated yet");
153
154        // SAFETY: `layout` has a non-zero size, and the block is freed below
155        // through the same allocator and layout it was allocated with.
156        let ptr = unsafe { alloc.alloc(layout) };
157        assert!(!ptr.is_null(), "system allocator returned null for 4 KiB");
158        // SAFETY: `ptr` is a live block of `SIZE` bytes from the line above.
159        unsafe { ptr.write_bytes(0xAB, SIZE) };
160
161        let during = BLOCK.snapshot().expect("the wrapper just allocated");
162        assert_eq!(during.live_bytes, SIZE as u64);
163        assert_eq!(during.alloc_count, 1);
164        assert_eq!(during.free_count, 0);
165
166        // SAFETY: `ptr` came from `alloc.alloc(layout)` above and has not been
167        // freed or reallocated since.
168        unsafe { alloc.dealloc(ptr, layout) };
169        let after = BLOCK.snapshot().expect("counters are live");
170        assert_eq!(after.live_bytes, 0);
171        assert_eq!(after.alloc_count, 1);
172        assert_eq!(after.free_count, 1);
173        assert_eq!(after.peak_bytes, SIZE as u64);
174    }
175
176    // The zeroing path is the inner allocator's, and it counts like any other
177    // allocation.
178    #[test]
179    fn alloc_zeroed_forwards_the_zeroing_and_counts() {
180        static BLOCK: Counters = Counters::new();
181        let alloc = TrackingAlloc::with_counters(System, &BLOCK);
182        let layout = layout(SIZE);
183
184        // SAFETY: `layout` has a non-zero size, and the block is freed below
185        // through the same allocator and layout it was allocated with.
186        let ptr = unsafe { alloc.alloc_zeroed(layout) };
187        assert!(!ptr.is_null(), "system allocator returned null for 4 KiB");
188        // SAFETY: `ptr` is a live block of `SIZE` bytes the allocator just
189        // initialized, so a shared slice over it is valid.
190        let bytes = unsafe { core::slice::from_raw_parts(ptr, SIZE) };
191        assert!(bytes.iter().all(|&b| b == 0), "the block was not zeroed");
192
193        let during = BLOCK.snapshot().expect("the wrapper just allocated");
194        assert_eq!(during.live_bytes, SIZE as u64);
195        assert_eq!(during.alloc_count, 1);
196
197        // SAFETY: `ptr` came from `alloc.alloc_zeroed(layout)` above and has
198        // not been freed or reallocated since.
199        unsafe { alloc.dealloc(ptr, layout) };
200        assert_eq!(BLOCK.snapshot().expect("counters are live").live_bytes, 0);
201    }
202
203    // A resize moves live bytes by the delta in either direction, and is
204    // neither an allocation nor a free.
205    #[test]
206    fn realloc_counts_the_resize_only() {
207        static BLOCK: Counters = Counters::new();
208        let alloc = TrackingAlloc::with_counters(System, &BLOCK);
209
210        // SAFETY: the layout has a non-zero size, and the block is resized and
211        // freed below through the same allocator.
212        let ptr = unsafe { alloc.alloc(layout(SIZE)) };
213        assert!(!ptr.is_null(), "system allocator returned null for 4 KiB");
214
215        // SAFETY: `ptr` came from `alloc.alloc` with `layout(SIZE)`, and
216        // `GROWN` is a non-zero size valid for that layout's alignment.
217        let ptr = unsafe { alloc.realloc(ptr, layout(SIZE), GROWN) };
218        assert!(!ptr.is_null(), "system allocator returned null for 64 KiB");
219        let grown = BLOCK.snapshot().expect("the wrapper just allocated");
220        assert_eq!(grown.live_bytes, GROWN as u64);
221        assert_eq!(grown.alloc_count, 1, "a resize is not an allocation");
222        assert_eq!(grown.free_count, 0, "a resize is not a free");
223
224        // SAFETY: the resize above left `ptr` holding `GROWN` bytes at that
225        // layout's alignment, and `SHRUNK` is a non-zero size.
226        let ptr = unsafe { alloc.realloc(ptr, layout(GROWN), SHRUNK) };
227        assert!(!ptr.is_null(), "system allocator returned null for 1 KiB");
228        let shrunk = BLOCK.snapshot().expect("counters are live");
229        assert_eq!(shrunk.live_bytes, SHRUNK as u64);
230        assert_eq!(shrunk.peak_bytes, GROWN as u64);
231
232        // SAFETY: `ptr` holds `SHRUNK` bytes from the resize above.
233        unsafe { alloc.dealloc(ptr, layout(SHRUNK)) };
234        let after = BLOCK.snapshot().expect("counters are live");
235        assert_eq!(after.live_bytes, 0);
236        assert_eq!(after.alloc_count, 1);
237        assert_eq!(after.free_count, 1);
238    }
239
240    // A null return is a block that was never handed out. Counting it would
241    // leave live bytes nothing can ever free.
242    #[test]
243    fn a_failed_allocation_counts_nothing() {
244        static BLOCK: Counters = Counters::new();
245        let alloc = TrackingAlloc::with_counters(Failing, &BLOCK);
246
247        // SAFETY: the inner allocator returns null without allocating, so
248        // there is no block to free.
249        assert!(unsafe { alloc.alloc(layout(SIZE)) }.is_null());
250        // SAFETY: as above.
251        assert!(unsafe { alloc.alloc_zeroed(layout(SIZE)) }.is_null());
252        assert_eq!(BLOCK.snapshot(), None, "a failed allocation was counted");
253    }
254}
concinnity_core/memory/tracking.rs

concinnity_core/memory/
tracking.rs
