fluxbench_core/
allocator.rs

1//! Global Allocator Interceptor
2//!
3//! Tracks heap allocations during benchmark execution.
4
5use std::alloc::{GlobalAlloc, Layout, System};
6use std::sync::atomic::{AtomicU64, Ordering};
7
8/// Process-wide allocation counters.
9///
10/// These are global atomics, not thread-local. Counts include allocations from
11/// every thread, so multi-thread async benchmarks include other threads' allocations.
12static ALLOCATED_BYTES: AtomicU64 = AtomicU64::new(0);
13static ALLOCATION_COUNT: AtomicU64 = AtomicU64::new(0);
14
15/// Tracking allocator that wraps the system allocator
16pub struct TrackingAllocator;
17
18unsafe impl GlobalAlloc for TrackingAllocator {
19    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
20        // SAFETY: Delegates to the system allocator with the provided layout.
21        let ptr = unsafe { System.alloc(layout) };
22        if !ptr.is_null() {
23            ALLOCATED_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed);
24            ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed);
25        }
26        ptr
27    }
28
29    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
30        // SAFETY: Pointer/layout are provided by corresponding allocation operations.
31        unsafe { System.dealloc(ptr, layout) };
32        // Note: We don't decrement on dealloc to track total allocations
33    }
34
35    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
36        // SAFETY: Delegates to the system allocator with the provided layout.
37        let ptr = unsafe { System.alloc_zeroed(layout) };
38        if !ptr.is_null() {
39            ALLOCATED_BYTES.fetch_add(layout.size() as u64, Ordering::Relaxed);
40            ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed);
41        }
42        ptr
43    }
44
45    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
46        // SAFETY: Pointer/layout originate from allocator contracts; new_size is caller-provided.
47        let new_ptr = unsafe { System.realloc(ptr, layout, new_size) };
48        if !new_ptr.is_null() {
49            // Track the size difference
50            if new_size > layout.size() {
51                ALLOCATED_BYTES.fetch_add((new_size - layout.size()) as u64, Ordering::Relaxed);
52            }
53            ALLOCATION_COUNT.fetch_add(1, Ordering::Relaxed);
54        }
55        new_ptr
56    }
57}
58
59/// Get current allocation statistics
60pub fn current_allocation() -> (u64, u64) {
61    (
62        ALLOCATED_BYTES.load(Ordering::Relaxed),
63        ALLOCATION_COUNT.load(Ordering::Relaxed),
64    )
65}
66
67/// Reset allocation counters (call before each iteration)
68pub fn reset_allocation_counter() {
69    ALLOCATED_BYTES.store(0, Ordering::Relaxed);
70    ALLOCATION_COUNT.store(0, Ordering::Relaxed);
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::sync::Mutex;
77
78    /// The counters are process-wide. Tests that write them take this lock.
79    static COUNTER_LOCK: Mutex<()> = Mutex::new(());
80
81    fn lock_counters() -> std::sync::MutexGuard<'static, ()> {
82        COUNTER_LOCK
83            .lock()
84            .unwrap_or_else(std::sync::PoisonError::into_inner)
85    }
86
87    #[test]
88    fn test_reset_allocation_counter() {
89        let _guard = lock_counters();
90        // Set some values
91        ALLOCATED_BYTES.store(1000, std::sync::atomic::Ordering::Relaxed);
92        ALLOCATION_COUNT.store(5, std::sync::atomic::Ordering::Relaxed);
93
94        // Reset
95        reset_allocation_counter();
96
97        let (bytes, count) = current_allocation();
98        assert_eq!(bytes, 0);
99        assert_eq!(count, 0);
100    }
101
102    #[test]
103    fn test_current_allocation_reads_atomics() {
104        let _guard = lock_counters();
105        ALLOCATED_BYTES.store(2048, std::sync::atomic::Ordering::Relaxed);
106        ALLOCATION_COUNT.store(10, std::sync::atomic::Ordering::Relaxed);
107
108        let (bytes, count) = current_allocation();
109        assert_eq!(bytes, 2048);
110        assert_eq!(count, 10);
111
112        // Clean up
113        reset_allocation_counter();
114    }
115
116    // Note: Full allocation tracking test requires TrackingAllocator
117    // to be installed as #[global_allocator] in the binary crate.
118    // This is done by end users, not the library.
119}
fluxbench_core/allocator.rs

fluxbench_core/
allocator.rs
