OXYGENHUB
◆ v0.2.0 oxygen / crates / oxy-core / src / alloc_stats.rs
2.4 KiB · Rust · last changed Memory per phase with a counting allocator 1 month ago source blame history raw
1 
//! A counting allocator for measuring memory (`--features alloc-stats` on
2 
//! the `oxy` binary, with `OXY_TIMINGS=1`): each timing phase then also
3 
//! reports live bytes, the peak so far and what the phase allocated, so a
4 
//! profile names the structure that holds the memory.
5
6 
use std::alloc::{GlobalAlloc, Layout, System};
7 
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed};
8
9 
static LIVE: AtomicUsize = AtomicUsize::new(0);
10 
static PEAK: AtomicUsize = AtomicUsize::new(0);
11 
static TOTAL: AtomicUsize = AtomicUsize::new(0);
12 
static COUNT: AtomicUsize = AtomicUsize::new(0);
13 
static ACTIVE: AtomicBool = AtomicBool::new(false);
14
15 
/// The system allocator, counting. Install it with
16 
/// `#[global_allocator] static A: Counting = Counting;`.
17 
pub struct Counting;
18
19 
impl Counting {
20 
    fn grew(n: usize) {
21 
        ACTIVE.store(true, Relaxed);
22 
        let live = LIVE.fetch_add(n, Relaxed) + n;
23 
        TOTAL.fetch_add(n, Relaxed);
24 
        COUNT.fetch_add(1, Relaxed);
25 
        PEAK.fetch_max(live, Relaxed);
26 
    }
27 
}
28
29 
unsafe impl GlobalAlloc for Counting {
30 
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
31 
        let p = System.alloc(l);
32 
        if !p.is_null() {
33 
            Self::grew(l.size());
34 
        }
35 
        p
36 
    }
37 
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
38 
        let p = System.alloc_zeroed(l);
39 
        if !p.is_null() {
40 
            Self::grew(l.size());
41 
        }
42 
        p
43 
    }
44 
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
45 
        System.dealloc(p, l);
46 
        LIVE.fetch_sub(l.size(), Relaxed);
47 
    }
48 
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
49 
        let q = System.realloc(p, l, new);
50 
        if !q.is_null() {
51 
            LIVE.fetch_sub(l.size(), Relaxed);
52 
            Self::grew(new);
53 
        }
54 
        q
55 
    }
56 
}
57
58 
/// What the counting allocator has seen, when it is installed.
59 
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
60 
pub struct Stats {
61 
    /// Bytes allocated and not yet freed.
62 
    pub live: usize,
63 
    /// The most `live` has been.
64 
    pub peak: usize,
65 
    /// Bytes ever allocated.
66 
    pub total: usize,
67 
    /// Allocations ever made.
68 
    pub count: usize,
69 
}
70
71 
/// The counters, or `None` when the counting allocator is not installed.
72 
pub fn stats() -> Option<Stats> {
73 
    ACTIVE.load(Relaxed).then(|| Stats {
74 
        live: LIVE.load(Relaxed),
75 
        peak: PEAK.load(Relaxed),
76 
        total: TOTAL.load(Relaxed),
77 
        count: COUNT.load(Relaxed),
78 
    })
79 
}