concinnity_memory/
tracking.rs

1// concinnity-memory/src/alloc.rs
2//
3// The `GlobalAlloc` wrapper that counts what the Rust heap is holding.
4//
5// Installing it is the binary's call, since `#[global_allocator]` is a
6// per-program item and an embedder may want an allocator of its own:
7//
8//     #[global_allocator]
9//     static ALLOC: TrackingAlloc<System> = TrackingAlloc::new(System);
10//
11// Until some binary does, `stats()` reports `None` rather than a zero that would
12// read as "the heap is empty".
13
14use core::alloc::{GlobalAlloc, Layout};
15
16use crate::counters::Counters;
17
18pub(crate) static COUNTERS: Counters = Counters::new();
19
20/// A `GlobalAlloc` that counts allocations and forwards them to `A`.
21pub struct TrackingAlloc<A> {
22    inner: A,
23}
24
25impl<A> TrackingAlloc<A> {
26    /// Wrap `inner` so every allocation through it is counted.
27    pub const fn new(inner: A) -> Self {
28        Self { inner }
29    }
30}
31
32// SAFETY: every method forwards to `inner`, which upholds the `GlobalAlloc`
33// contract, and returns its pointer unchanged. The counter updates are relaxed
34// atomics over statically allocated storage, so the allocator cannot re-enter
35// itself.
36unsafe impl<A: GlobalAlloc> GlobalAlloc for TrackingAlloc<A> {
37    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
38        // SAFETY: `layout` is forwarded unchanged from our caller, who upholds
39        // the `GlobalAlloc::alloc` contract.
40        let ptr = unsafe { self.inner.alloc(layout) };
41        if !ptr.is_null() {
42            COUNTERS.record_alloc(layout.size());
43            crate::detail::record_alloc(layout.size());
44        }
45        ptr
46    }
47
48    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
49        // SAFETY: as `alloc`, forwarding our caller's obligations to `inner`.
50        let ptr = unsafe { self.inner.alloc_zeroed(layout) };
51        if !ptr.is_null() {
52            COUNTERS.record_alloc(layout.size());
53            crate::detail::record_alloc(layout.size());
54        }
55        ptr
56    }
57
58    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
59        // SAFETY: `ptr` and `layout` come from our caller, who guarantees the
60        // block was allocated by us with this layout; we allocated it through
61        // `inner`, so `inner` is the right allocator to free it.
62        unsafe { self.inner.dealloc(ptr, layout) };
63        COUNTERS.record_free(layout.size());
64        crate::detail::record_free(layout.size());
65    }
66
67    // Forwarded rather than left to the default alloc-copy-dealloc so `inner`
68    // can grow a block in place. A failed resize leaves the original block
69    // allocated and unchanged, so the counters only move on success.
70    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
71        // SAFETY: `ptr`, `layout` and `new_size` are forwarded unchanged from
72        // our caller, who upholds the `GlobalAlloc::realloc` contract.
73        let new_ptr = unsafe { self.inner.realloc(ptr, layout, new_size) };
74        if !new_ptr.is_null() {
75            COUNTERS.record_realloc(layout.size(), new_size);
76            crate::detail::record_realloc(layout.size(), new_size);
77        }
78        new_ptr
79    }
80}
81
82/// Install [`TrackingAlloc`] over the system allocator as this program's
83/// `#[global_allocator]`.
84///
85/// `#[global_allocator]` is a per-program item, so every binary, test binary,
86/// and benchmark that should count its own heap invokes this once at its crate
87/// root. A program without it runs on Rust's default allocator and reports no
88/// memory at all, which is what [`crate::stats`] returning `None` means.
89///
90/// ```
91/// concinnity_memory::install_global_allocator!();
92/// ```
93#[macro_export]
94macro_rules! install_global_allocator {
95    () => {
96        #[global_allocator]
97        static CN_GLOBAL_ALLOC: $crate::TrackingAlloc<std::alloc::System> =
98            $crate::TrackingAlloc::new(std::alloc::System);
99    };
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use std::alloc::System;
106
107    // Drives the real wrapper against the real system allocator. This is the
108    // one test that touches the process-global counters, so it asserts on
109    // deltas rather than absolute values.
110    #[test]
111    fn tracking_alloc_forwards_and_counts() {
112        let alloc = TrackingAlloc::new(System);
113        let layout = Layout::from_size_align(4096, 8).expect("valid layout");
114
115        let before = crate::stats().map(|s| s.live_bytes).unwrap_or(0);
116        // SAFETY: `layout` has a non-zero size, and the block is freed below
117        // through the same allocator and layout it was allocated with.
118        let ptr = unsafe { alloc.alloc(layout) };
119        assert!(!ptr.is_null(), "system allocator returned null for 4 KiB");
120
121        let during = crate::stats()
122            .expect("the wrapper just allocated")
123            .live_bytes;
124        assert!(
125            during >= before + 4096,
126            "live bytes {during} did not rise by the allocation size from {before}"
127        );
128
129        // SAFETY: `ptr` came from `alloc.alloc(layout)` above and has not been
130        // freed or reallocated since.
131        unsafe { alloc.dealloc(ptr, layout) };
132        assert!(crate::stats().expect("counters are live").live_bytes >= before);
133    }
134}
concinnity_memory/tracking.rs

concinnity_memory/
tracking.rs
