alloc_metrics/
lib.rs

1#![no_std]
2#![doc = include_str!("../README.md")]
3
4#[cfg(feature = "std")]
5extern crate std;
6
7use core::alloc::GlobalAlloc;
8
9use derive_more::{Add, AddAssign, Sub, SubAssign, Neg};
10
11#[cfg(test)]
12mod test;
13
14#[cfg(feature = "thread")]
15std::thread_local! {
16    static THREAD_METRICS: core::cell::Cell<Metrics> = core::cell::Cell::new(Metrics { allocated_bytes: 0, allocations: 0 });
17}
18#[cfg(feature = "global")]
19static GLOBAL_METRICS: spin::Mutex<Metrics> = spin::Mutex::new(Metrics { allocated_bytes: 0, allocations: 0 });
20
21fn add_assign_metrics(_metrics: Metrics) {
22    #[cfg(feature = "thread")]
23    { THREAD_METRICS.set(THREAD_METRICS.get() + _metrics) }
24    #[cfg(feature = "global")]
25    { *GLOBAL_METRICS.lock() += _metrics }
26}
27
28/// Holds metrics on memory allocations.
29/// 
30/// This type implements several arithmetic operations, which allows for computing metric deltas and aggregating them.
31#[cfg_attr(feature = "thread", doc = r#"
32```rust
33let before = alloc_metrics::thread_metrics();
34// ... do some work ...
35let after = alloc_metrics::thread_metrics();
36let delta = after - before;
37```
38"#)]
39#[derive(Clone, Copy, Debug, PartialEq, Eq, Add, AddAssign, Sub, SubAssign, Neg)]
40pub struct Metrics {
41    /// The number of bytes that were allocated.
42    pub allocated_bytes: isize,
43    /// The number of allocations made.
44    pub allocations: isize,
45}
46
47/// A global allocator type that tracks allocation metrics.
48/// 
49/// This type makes use of shared memory in order to aggregate metrics while still supporting arbitrary global allocator composition.
50/// Because of this, there should only ever be at most one instance of it (hence being a global allocator).
51pub struct MetricAlloc<A: GlobalAlloc> {
52    wrapped: A,
53}
54impl<A: GlobalAlloc> MetricAlloc<A> {
55    /// Wraps an existing global allocator into a metric allocator.
56    pub const fn new(wrapped: A) -> Self {
57        Self { wrapped }
58    }
59}
60unsafe impl<A: GlobalAlloc> GlobalAlloc for MetricAlloc<A> {
61    unsafe fn alloc(&self, layout: core::alloc::Layout) -> *mut u8 {
62        let res = self.wrapped.alloc(layout);
63        if !res.is_null() {
64            add_assign_metrics(Metrics {
65                allocated_bytes: layout.size() as isize,
66                allocations: 1,
67            });
68        }
69        res
70    }
71    unsafe fn alloc_zeroed(&self, layout: core::alloc::Layout) -> *mut u8 {
72        let res = self.wrapped.alloc_zeroed(layout);
73        if !res.is_null() {
74            add_assign_metrics(Metrics {
75                allocated_bytes: layout.size() as isize,
76                allocations: 1,
77            });
78        }
79        res
80    }
81    unsafe fn dealloc(&self, ptr: *mut u8, layout: core::alloc::Layout) {
82        self.wrapped.dealloc(ptr, layout);
83        if !ptr.is_null() {
84            add_assign_metrics(Metrics {
85                allocated_bytes: -(layout.size() as isize),
86                allocations: -1,
87            });
88        }
89    }
90    unsafe fn realloc(&self, ptr: *mut u8, layout: core::alloc::Layout, new_size: usize) -> *mut u8 {
91        let res = self.wrapped.realloc(ptr, layout, new_size);
92        if !res.is_null() {
93            add_assign_metrics(Metrics {
94                allocated_bytes: new_size as isize - layout.size() as isize,
95                allocations: 0,
96            })
97        }
98        res
99    }
100}
101
102/// Get the current allocation metrics for the current thread.
103/// 
104/// Allocations in other threads will not affect this value.
105#[cfg(feature = "thread")]
106pub fn thread_metrics() -> Metrics {
107    THREAD_METRICS.get()
108}
109
110/// Get the current allocation metrics for the entire program.
111/// 
112/// Allocations in other threads will affect this value, but each modification is guaranteed to be atomic.
113#[cfg(feature = "global")]
114pub fn global_metrics() -> Metrics {
115    *GLOBAL_METRICS.lock()
116}
alloc_metrics/lib.rs

alloc_metrics/
lib.rs
