metricus_allocator/
lib.rs

1#![doc = include_str!("../README.md")]
2
3use metricus::{Counter, CounterOps, Id, PreAllocatedMetric};
4use std::alloc::{GlobalAlloc, Layout};
5use std::cell::Cell;
6use std::sync::LazyLock;
7
8const ALLOC_COUNTER_ID: Id = Id::MAX - 1004;
9const ALLOC_BYTES_COUNTER_ID: Id = Id::MAX - 1003;
10const DEALLOC_COUNTER_ID: Id = Id::MAX - 1002;
11const DEALLOC_BYTES_COUNTER_ID: Id = Id::MAX - 1001;
12
13const fn get_aligned_size(layout: Layout) -> usize {
14    let alignment_mask: usize = layout.align() - 1;
15    (layout.size() + alignment_mask) & !alignment_mask
16}
17
18/// This allocator will use instrumentation to count the number of allocations and de-allocations
19/// occurring in the program. All calls to allocate (and free) memory are delegated to the concrete
20/// allocator (`std::alloc::System` by default). Once the allocator has been registered as
21/// `global_allocator` you need to call [enable_allocator_instrumentation] from each thread that
22/// wants to include its allocation and de-allocation metrics.
23///
24/// ```no_run
25/// use metricus_allocator::CountingAllocator;
26///
27/// #[global_allocator]
28/// static GLOBAL: CountingAllocator = CountingAllocator;
29/// ```
30pub struct CountingAllocator;
31
32#[allow(static_mut_refs)]
33unsafe impl GlobalAlloc for CountingAllocator {
34    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
35        // provide metrics only if instrumentation has been enabled for this thread
36        if INSTRUMENTATION_ENABLED.get() {
37            COUNTERS.alloc_count.increment();
38            COUNTERS.alloc_bytes.increment_by(get_aligned_size(layout) as u64);
39        }
40
41        // delegate to the appropriate allocator
42        #[cfg(all(feature = "jemalloc", not(feature = "mimalloc")))]
43        {
44            return unsafe { jemallocator::Jemalloc.alloc(layout) };
45        }
46        #[cfg(all(feature = "mimalloc", not(feature = "jemalloc")))]
47        {
48            return unsafe { mimalloc::MiMalloc.alloc(layout) };
49        }
50        unsafe { std::alloc::System.alloc(layout) }
51    }
52
53    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
54        // provide metrics only if instrumentation has been enabled for this thread
55        if INSTRUMENTATION_ENABLED.get() {
56            COUNTERS.dealloc_count.increment();
57            COUNTERS.dealloc_bytes.increment_by(get_aligned_size(layout) as u64);
58        }
59
60        // delegate to the appropriate allocator
61        #[cfg(all(feature = "jemalloc", not(feature = "mimalloc")))]
62        {
63            unsafe {
64                jemallocator::Jemalloc.dealloc(ptr, layout);
65            }
66            return;
67        }
68        #[cfg(all(feature = "mimalloc", not(feature = "jemalloc")))]
69        {
70            unsafe {
71                mimalloc::MiMalloc.dealloc(ptr, layout);
72            }
73            return;
74        }
75        unsafe { std::alloc::System.dealloc(ptr, layout) }
76    }
77}
78
79impl CountingAllocator {
80    /// Default counters to be used with the `CountingAllocator`.
81    pub fn metrics() -> Vec<PreAllocatedMetric> {
82        vec![
83            PreAllocatedMetric::counter("global_allocator", ALLOC_COUNTER_ID, &[("fn_name", "alloc")]),
84            PreAllocatedMetric::counter("global_allocator", ALLOC_BYTES_COUNTER_ID, &[("fn_name", "alloc_bytes")]),
85            PreAllocatedMetric::counter("global_allocator", DEALLOC_COUNTER_ID, &[("fn_name", "dealloc")]),
86            PreAllocatedMetric::counter("global_allocator", DEALLOC_BYTES_COUNTER_ID, &[("fn_name", "dealloc_bytes")]),
87        ]
88    }
89}
90
91thread_local! {
92    static INSTRUMENTATION_ENABLED: Cell<bool> = const { Cell::new(false) };
93}
94
95/// This should be called by a thread that wants to opt in to send allocation and de-allocation
96/// metrics. By default, per thread instrumentation is disabled. This is usually backend dependent
97/// as some backends can support sending metrics from multiple threads whereas others can be limited
98/// to the main thread only.
99///
100/// ## Examples
101///
102/// Enable instrumentation for the main thread.
103/// ```no_run
104///
105/// use metricus_allocator::enable_allocator_instrumentation;
106/// use metricus_allocator::CountingAllocator;
107///
108/// #[global_allocator]
109/// static GLOBAL: CountingAllocator = CountingAllocator;
110///
111/// fn main() {
112///     enable_allocator_instrumentation();
113/// }
114/// ```
115///
116/// Enable instrumentation for the background thread.
117/// ```no_run
118///
119/// use metricus_allocator::enable_allocator_instrumentation;
120/// use metricus_allocator::CountingAllocator;
121///
122/// #[global_allocator]
123/// static GLOBAL: CountingAllocator = CountingAllocator;
124///
125/// fn main() {
126///     let _ = std::thread::spawn(|| {
127///         enable_allocator_instrumentation();
128///     });
129/// }
130/// ```
131pub fn enable_allocator_instrumentation() {
132    INSTRUMENTATION_ENABLED.set(true);
133}
134
135static COUNTERS: LazyLock<Counters> = LazyLock::new(|| Counters {
136    // `counter_with_id` creates a counter object without registering it.
137    // These allocation counters are created lazily on first use and cache the active metrics handle.
138    // If they are initialized before `set_metrics`, they will remain bound to the no-op backend.
139    // Ensure the backend is set before enabling allocator instrumentation if you want these to emit.
140    alloc_count: Counter::new_with_id(ALLOC_COUNTER_ID),
141    alloc_bytes: Counter::new_with_id(ALLOC_BYTES_COUNTER_ID),
142    dealloc_count: Counter::new_with_id(DEALLOC_COUNTER_ID),
143    dealloc_bytes: Counter::new_with_id(DEALLOC_BYTES_COUNTER_ID),
144});
145
146struct Counters {
147    alloc_count: Counter,
148    alloc_bytes: Counter,
149    dealloc_count: Counter,
150    dealloc_bytes: Counter,
151}
metricus_allocator/lib.rs

metricus_allocator/
lib.rs
