forge/
allocation.rs

1//! Allocation counting for Forge benchmark samples.
2
3use std::{
4    alloc::{GlobalAlloc, Layout},
5    sync::atomic::{AtomicBool, AtomicU64, Ordering},
6};
7
8use serde::{Deserialize, Serialize};
9
10/// Allocation totals captured while counting is enabled.
11#[derive(Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
12pub(crate) struct AllocationSnapshot {
13    pub count: u64,
14    pub bytes: u64,
15}
16
17/// Transparent counting wrapper around Forge's existing allocator.
18pub(crate) struct CountingAllocator<A> {
19    inner: A,
20    enabled: AtomicBool,
21    count: AtomicU64,
22    bytes: AtomicU64,
23}
24
25impl<A> CountingAllocator<A> {
26    pub(crate) const fn new(inner: A) -> Self {
27        Self {
28            inner,
29            enabled: AtomicBool::new(false),
30            count: AtomicU64::new(0),
31            bytes: AtomicU64::new(0),
32        }
33    }
34
35    pub(crate) fn start(&self) {
36        self.enabled.store(false, Ordering::SeqCst);
37        self.count.store(0, Ordering::Relaxed);
38        self.bytes.store(0, Ordering::Relaxed);
39        self.enabled.store(true, Ordering::SeqCst);
40    }
41
42    pub(crate) fn stop(&self) -> AllocationSnapshot {
43        self.enabled.store(false, Ordering::SeqCst);
44
45        AllocationSnapshot {
46            count: self.count.load(Ordering::Relaxed),
47            bytes: self.bytes.load(Ordering::Relaxed),
48        }
49    }
50
51    #[inline]
52    fn record(&self, size: usize) {
53        if !self.enabled.load(Ordering::Relaxed) {
54            return;
55        }
56
57        self.count.fetch_add(1, Ordering::Relaxed);
58        self.bytes
59            .fetch_add(u64::try_from(size).unwrap_or(u64::MAX), Ordering::Relaxed);
60    }
61}
62
63// SAFETY: every operation delegates unchanged layouts and pointers to the
64// wrapped global allocator. The wrapper observes successful allocation
65// requests without changing allocation or deallocation semantics.
66unsafe impl<A> GlobalAlloc for CountingAllocator<A>
67where
68    A: GlobalAlloc,
69{
70    // SAFETY: the caller supplies a valid allocation layout under the `GlobalAlloc` contract.
71    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
72        // SAFETY: upheld by the caller of GlobalAlloc::alloc.
73        let pointer = unsafe { self.inner.alloc(layout) };
74
75        if !pointer.is_null() {
76            self.record(layout.size());
77        }
78
79        pointer
80    }
81
82    // SAFETY: the caller supplies a valid allocation layout under the `GlobalAlloc` contract.
83    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
84        // SAFETY: upheld by the caller of GlobalAlloc::alloc_zeroed.
85        let pointer = unsafe { self.inner.alloc_zeroed(layout) };
86
87        if !pointer.is_null() {
88            self.record(layout.size());
89        }
90
91        pointer
92    }
93
94    // SAFETY: the pointer and layout come from this allocator under the `GlobalAlloc` contract.
95    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
96        // SAFETY: upheld by the caller of GlobalAlloc::dealloc.
97        unsafe { self.inner.dealloc(ptr, layout) };
98    }
99
100    // SAFETY: the pointer, old layout, and new size satisfy the `GlobalAlloc` contract.
101    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
102        // SAFETY: upheld by the caller of GlobalAlloc::realloc.
103        let new_pointer = unsafe { self.inner.realloc(ptr, layout, new_size) };
104
105        if !new_pointer.is_null() {
106            self.record(new_size);
107        }
108
109        new_pointer
110    }
111}
112
113#[cfg(test)]
114mod tests {
115    use std::alloc::{GlobalAlloc as _, System};
116
117    use googletest::prelude::*;
118
119    use super::*;
120
121    #[gtest]
122    fn counting_wrapper_records_successful_allocation_requests() -> Result<()> {
123        let allocator = CountingAllocator::new(System);
124        let layout = Layout::from_size_align(64, 8).or_fail()?;
125
126        allocator.start();
127
128        // SAFETY: `layout` is valid and the returned pointer is deallocated below with the
129        // same allocator and layout.
130        let pointer = unsafe { allocator.alloc(layout) };
131
132        verify_false!(pointer.is_null())?;
133        verify_that!(
134            allocator.stop(),
135            eq(AllocationSnapshot {
136                count: 1,
137                bytes: 64
138            })
139        )?;
140
141        // SAFETY: `pointer` was allocated above by `allocator` with this exact layout.
142        unsafe { allocator.dealloc(pointer, layout) };
143
144        Ok(())
145    }
146}
forge/allocation.rs

forge/
allocation.rs
