alloc_test/alloc/
measure.rs

1use std::{
2    mem,
3    sync::atomic::{AtomicBool, Ordering},
4};
5
6use derive_more::Display;
7use serde::{Deserialize, Serialize};
8
9#[derive(Debug, Default, Clone, Display, Serialize, Deserialize)]
10#[display(fmt = r#"Currently allocated (B): {current}
11Maximum allocated (B): {peak}
12Total amount of claimed memory (B): {total_size}
13Total number of allocations: (N): {total_num}
14Reallocations (N): {reallocs}
15"#)]
16pub struct MemoryStats {
17    pub current: usize,
18    pub peak: usize,
19    pub total_size: usize,
20    pub total_num: usize,
21    pub reallocs: usize,
22}
23
24static mut TRACE_ALLOCS: AtomicBool = AtomicBool::new(false);
25
26static mut ALLOC_STATS: MemoryStats = MemoryStats {
27    current: 0,
28    peak: 0,
29    total_size: 0,
30    total_num: 0,
31    reallocs: 0,
32};
33
34/// Traces allocations performed while executing the `f`.
35///
36/// Beware that allocations made by nother threads will be also recorded.
37///
38/// ```ignore
39/// use alloc_test::{TracingAllocator, MemoryTracingHooks, trace_allocs};
40/// use std::alloc::System;
41///
42/// #[global_allocator]
43/// static ALLOCATOR: TracingAllocator<MemoryTracingHooks, System> =
44///     TracingAllocator::new(MemoryTracingHooks, System);
45///
46/// fn main() {
47///     let (_, stats) = trace_allocs(|| {
48///         let r: Vec<u8> = vec![1, 2, 3];
49///         r
50///     });
51///     assert_eq!(stats.peak, 3);
52///     assert_eq!(stats.current, 3);
53/// }
54/// ```
55#[allow(static_mut_refs)]
56pub fn trace_allocs<F: FnOnce() -> O, O>(f: F) -> (O, MemoryStats) {
57    unsafe {
58        while TRACE_ALLOCS
59            .compare_exchange(false, true, Ordering::Acquire, Ordering::Acquire)
60            .is_err()
61        {}
62        let o = f();
63        let stats = mem::take(&mut ALLOC_STATS);
64        TRACE_ALLOCS.store(false, Ordering::Release);
65        (o, stats)
66    }
67}
68
69pub struct MemoryTracingHooks;
70
71#[allow(static_mut_refs)]
72unsafe impl super::allocator::AllocHooks for MemoryTracingHooks {
73    fn on_alloc(&self, _pointer: *mut u8, size: usize, _align: usize) {
74        unsafe {
75            if TRACE_ALLOCS.load(Ordering::Acquire) {
76                // println!("allocating {size}");
77                ALLOC_STATS.current += size;
78                ALLOC_STATS.total_size += size;
79                ALLOC_STATS.total_num += 1;
80                if ALLOC_STATS.current > ALLOC_STATS.peak {
81                    ALLOC_STATS.peak = ALLOC_STATS.current;
82                }
83            }
84        }
85    }
86
87    fn on_dealloc(&self, _pointer: *mut u8, size: usize, _align: usize) {
88        unsafe {
89            if TRACE_ALLOCS.load(Ordering::Acquire) {
90                ALLOC_STATS.current = ALLOC_STATS.current.saturating_sub(size);
91            }
92        }
93    }
94
95    fn on_alloc_zeroed(&self, pointer: *mut u8, size: usize, align: usize) {
96        self.on_alloc(pointer, size, align);
97    }
98
99    fn on_realloc(
100        &self,
101        old_pointer: *mut u8,
102        new_pointer: *mut u8,
103        old_size: usize,
104        new_size: usize,
105        align: usize,
106    ) {
107        unsafe {
108            if TRACE_ALLOCS.load(Ordering::Acquire) {
109                // println!("reallocating {old_size} -> {new_size}");
110                ALLOC_STATS.reallocs += 1;
111            }
112        }
113        self.on_dealloc(old_pointer, old_size, align);
114        self.on_alloc(new_pointer, new_size, align);
115    }
116}
alloc_test/alloc/measure.rs

alloc_test/alloc/
measure.rs
