simoom_sys/
alloc.rs

1//! A counting global allocator for memory measurements in tests and benchmarks.
2//!
3//! ```ignore
4//! #[global_allocator]
5//! static A: simoom_sys::alloc::Counting = simoom_sys::alloc::Counting;
6//! ```
7//!
8//! It forwards to the system allocator and keeps the net number of live heap bytes, in total
9//! and per thread (bytes freed on another thread than the one that allocated them count against
10//! the freeing thread).
11
12use std::alloc::{GlobalAlloc, Layout, System};
13use std::cell::Cell;
14use std::sync::atomic::{AtomicI64, Ordering};
15
16/// The counting allocator. See the [module docs](self).
17#[derive(Debug, Clone, Copy, Default)]
18pub struct Counting;
19
20static TOTAL: AtomicI64 = AtomicI64::new(0);
21/// Calls to `alloc`, `alloc_zeroed` and `realloc`.
22static CALLS: AtomicI64 = AtomicI64::new(0);
23/// Allocations not freed yet.
24static LIVE: AtomicI64 = AtomicI64::new(0);
25
26thread_local! {
27    static LOCAL: Cell<i64> = const { Cell::new(0) };
28    /// The highest value `LOCAL` reached since the last [`reset_thread_peak`].
29    static PEAK: Cell<i64> = const { Cell::new(0) };
30}
31
32fn add(n: i64) {
33    TOTAL.fetch_add(n, Ordering::Relaxed);
34    if n > 0 {
35        CALLS.fetch_add(1, Ordering::Relaxed);
36    }
37    let _ = LOCAL.try_with(|c| {
38        let v = c.get() + n;
39        c.set(v);
40        let _ = PEAK.try_with(|p| p.set(p.get().max(v)));
41    });
42}
43
44// SAFETY: every method forwards to `System` with the same arguments and only updates counters,
45// which never allocate (the thread-local is const-initialised and has no destructor).
46unsafe impl GlobalAlloc for Counting {
47    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
48        // SAFETY: forwarded unchanged; the caller upholds `GlobalAlloc::alloc`'s contract.
49        let p = unsafe { System.alloc(layout) };
50        if !p.is_null() {
51            add(layout.size() as i64);
52            LIVE.fetch_add(1, Ordering::Relaxed);
53        }
54        p
55    }
56
57    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
58        // SAFETY: forwarded unchanged.
59        let p = unsafe { System.alloc_zeroed(layout) };
60        if !p.is_null() {
61            add(layout.size() as i64);
62            LIVE.fetch_add(1, Ordering::Relaxed);
63        }
64        p
65    }
66
67    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
68        // SAFETY: forwarded unchanged; `ptr` came from this allocator, so from `System`.
69        unsafe { System.dealloc(ptr, layout) };
70        add(-(layout.size() as i64));
71        LIVE.fetch_sub(1, Ordering::Relaxed);
72    }
73
74    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
75        // SAFETY: forwarded unchanged.
76        let p = unsafe { System.realloc(ptr, layout, new_size) };
77        if !p.is_null() {
78            add(new_size as i64 - layout.size() as i64);
79        }
80        p
81    }
82}
83
84/// Net live heap bytes of the whole process (when [`Counting`] is the global allocator).
85pub fn total_bytes() -> i64 {
86    TOTAL.load(Ordering::Relaxed)
87}
88
89/// Heap allocations of the whole process not freed yet (when [`Counting`] is the global
90/// allocator).
91pub fn live_allocations() -> i64 {
92    LIVE.load(Ordering::Relaxed)
93}
94
95/// Allocation calls so far (allocations, and reallocations that grow), process-wide.
96pub fn allocation_calls() -> i64 {
97    CALLS.load(Ordering::Relaxed)
98}
99
100/// Net heap bytes allocated minus freed by the calling thread.
101pub fn thread_bytes() -> i64 {
102    LOCAL.try_with(|c| c.get()).unwrap_or(0)
103}
104
105/// The highest [`thread_bytes`] of the calling thread since its last [`reset_thread_peak`]
106/// (or since it started): what a piece of code needs at its peak, scratch included.
107pub fn thread_peak_bytes() -> i64 {
108    PEAK.try_with(|c| c.get()).unwrap_or(0)
109}
110
111/// Starts a new peak for [`thread_peak_bytes`] at the calling thread's current
112/// [`thread_bytes`].
113pub fn reset_thread_peak() {
114    let now = thread_bytes();
115    let _ = PEAK.try_with(|c| c.set(now));
116}
117
118#[cfg(all(target_os = "linux", target_env = "gnu"))]
119extern "C" {
120    fn malloc_trim(pad: usize) -> core::ffi::c_int;
121}
122
123/// Asks the C allocator to give memory it holds free back to the operating system (glibc's
124/// `malloc_trim`; elsewhere nothing). The event loops call it once after a burst, when they
125/// have been idle for a while (K2): glibc keeps what a thread freed in that thread's arena, so a
126/// burst of large responses otherwise left tens of MB resident.
127pub fn release_free_memory() {
128    #[cfg(all(target_os = "linux", target_env = "gnu"))]
129    // SAFETY: `malloc_trim` takes no pointers and is thread-safe (it locks each arena in turn).
130    unsafe {
131        malloc_trim(0);
132    }
133}
simoom_sys/alloc.rs

simoom_sys/
alloc.rs
