Skip to main content

subms_arena_allocator/features/
stats.rs

1//! `StatsBump`: bump arena with runtime counters for observability.
2//!
3//! Tracks:
4//! - `allocations` - total `alloc_*` calls served (excluding refusals).
5//! - `bytes_used` - sum of allocation `size` values (excludes padding).
6//! - `bytes_wasted` - sum of alignment padding inserted between
7//!   allocations.
8//! - `peak_bytes` - high-watermark of cursor across this arena's
9//!   lifetime.
10//! - `chunk_count` - number of distinct chunks the arena has ever
11//!   opened (grows monotonically; this struct also auto-grows so the
12//!   counter is meaningful).
13//!
14//! Counters survive `reset()` so a long-running process can read
15//! lifetime aggregates. `clear_stats()` zeros them.
16
17use std::alloc::{Layout, alloc, dealloc};
18use std::ptr;
19
20use crate::align_up;
21
22/// Snapshot of runtime counters.
23#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
24pub struct BumpStats {
25    pub allocations: u64,
26    pub bytes_used: u64,
27    pub bytes_wasted: u64,
28    pub peak_bytes: u64,
29    pub chunk_count: u64,
30}
31
32/// Bump arena with instrumentation. Auto-grows like `GrowableBump`.
33pub struct StatsBump {
34    chunks: Vec<Chunk>,
35    cursor: usize,
36    stats: BumpStats,
37}
38
39struct Chunk {
40    ptr: *mut u8,
41    layout: Layout,
42}
43
44impl Drop for Chunk {
45    fn drop(&mut self) {
46        unsafe { dealloc(self.ptr, self.layout) };
47    }
48}
49
50impl StatsBump {
51    /// New arena with a 4 KiB initial chunk.
52    pub fn new() -> Self {
53        Self::with_capacity(4096)
54    }
55
56    /// New arena with the requested initial chunk size.
57    pub fn with_capacity(initial: usize) -> Self {
58        let initial = initial.max(64);
59        let layout = Layout::from_size_align(initial, 16).expect("layout");
60        let ptr = unsafe { alloc(layout) };
61        assert!(!ptr.is_null(), "OOM allocating first stats chunk");
62        Self {
63            chunks: vec![Chunk { ptr, layout }],
64            cursor: 0,
65            stats: BumpStats {
66                chunk_count: 1,
67                ..Default::default()
68            },
69        }
70    }
71
72    /// Allocate a `Copy` value. Updates the counters.
73    pub fn alloc_copy<T: Copy>(&mut self, value: T) -> &mut T {
74        let layout = Layout::new::<T>();
75        let p = self.alloc_raw(layout);
76        unsafe {
77            ptr::write(p as *mut T, value);
78            &mut *(p as *mut T)
79        }
80    }
81
82    /// Allocate `layout.size()` bytes aligned to `layout.align()`.
83    /// Updates the counters.
84    pub fn alloc_raw(&mut self, layout: Layout) -> *mut u8 {
85        let size = layout.size();
86        let align = layout.align();
87        let last_size;
88        let last_ptr;
89        {
90            let last = self.chunks.last().expect("at least one chunk");
91            last_size = last.layout.size();
92            last_ptr = last.ptr;
93        }
94        let base = last_ptr as usize;
95        let aligned = align_up(base + self.cursor, align) - base;
96        let waste = aligned - self.cursor;
97        let end = aligned + size;
98        if end <= last_size {
99            self.cursor = end;
100            self.bump_counters(size as u64, waste as u64);
101            return unsafe { last_ptr.add(aligned) };
102        }
103        // Grow.
104        self.grow(size + align);
105        let last_ptr = self.chunks.last().unwrap().ptr;
106        let base = last_ptr as usize;
107        let aligned = align_up(base, align) - base;
108        // Fresh chunk starts at offset 0, so the waste counter only
109        // tracks the per-call padding inside the new chunk (typically
110        // zero because chunks are 16-byte aligned; non-zero only if
111        // the caller asks for >16-byte alignment).
112        self.cursor = aligned + size;
113        self.bump_counters(size as u64, aligned as u64);
114        unsafe { last_ptr.add(aligned) }
115    }
116
117    fn bump_counters(&mut self, size: u64, waste: u64) {
118        self.stats.allocations += 1;
119        self.stats.bytes_used += size;
120        self.stats.bytes_wasted += waste;
121        let cursor = self.cursor as u64;
122        if cursor > self.stats.peak_bytes {
123            self.stats.peak_bytes = cursor;
124        }
125    }
126
127    fn grow(&mut self, min_bytes: usize) {
128        let last = self.chunks.last().expect("at least one chunk");
129        let new_size = (last.layout.size() * 2).max(min_bytes);
130        let layout = Layout::from_size_align(new_size, 16).expect("layout");
131        let ptr = unsafe { alloc(layout) };
132        assert!(!ptr.is_null(), "OOM growing arena");
133        self.chunks.push(Chunk { ptr, layout });
134        self.cursor = 0;
135        self.stats.chunk_count += 1;
136    }
137
138    /// Rewind. Keeps the largest chunk. Stats are preserved.
139    pub fn reset(&mut self) {
140        if self.chunks.len() > 1 {
141            let largest = self
142                .chunks
143                .iter()
144                .enumerate()
145                .max_by_key(|(_, c)| c.layout.size())
146                .map(|(i, _)| i)
147                .unwrap();
148            let keeper = self.chunks.swap_remove(largest);
149            self.chunks.clear();
150            self.chunks.push(keeper);
151        }
152        self.cursor = 0;
153    }
154
155    /// Snapshot the live counters.
156    pub fn stats(&self) -> BumpStats {
157        self.stats
158    }
159
160    /// Zero the counters. The `chunk_count` is restored to the current
161    /// retained chunk count rather than zero, so the snapshot remains
162    /// meaningful immediately after.
163    pub fn clear_stats(&mut self) {
164        self.stats = BumpStats {
165            chunk_count: self.chunks.len() as u64,
166            ..Default::default()
167        };
168    }
169}
170
171impl Default for StatsBump {
172    fn default() -> Self {
173        Self::new()
174    }
175}
176
177#[cfg(test)]
178#[path = "stats_tests.rs"]
179mod tests;