rave_core/
debug_alloc.rs

1//! Debug-only host allocation tracker.
2//!
3//! Activated via `--features debug-alloc`.  Wraps the global allocator to
4//! count host-side heap allocations during steady-state processing.
5//!
6//! # Usage
7//!
8//! ```rust,ignore
9//! #[cfg(feature = "debug-alloc")]
10//! {
11//!     crate::debug_alloc::reset();
12//!     crate::debug_alloc::enable();
13//!     // ... inference stage runs one frame ...
14//!     crate::debug_alloc::disable();
15//!     let count = crate::debug_alloc::count();
16//!     assert_eq!(count, 0, "host allocations during inference: {count}");
17//! }
18//! ```
19//!
20//! Does NOT affect release builds (feature-gated, zero-cost when disabled).
21
22#[cfg(feature = "debug-alloc")]
23pub use inner::*;
24
25#[cfg(feature = "debug-alloc")]
26mod inner {
27    use std::alloc::{GlobalAlloc, Layout, System};
28    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
29
30    static TRACKING: AtomicBool = AtomicBool::new(false);
31    static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);
32
33    /// Global allocator wrapper that counts host allocations when enabled.
34    pub struct TrackingAllocator;
35
36    // SAFETY: delegates to `System` allocator for all actual work.
37    // The atomic counter adds negligible overhead (single relaxed fetch_add).
38    unsafe impl GlobalAlloc for TrackingAllocator {
39        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
40            if TRACKING.load(Ordering::Relaxed) {
41                ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
42            }
43            // SAFETY: delegates to System which upholds GlobalAlloc contract.
44            unsafe { System.alloc(layout) }
45        }
46
47        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
48            // SAFETY: ptr was allocated by System.alloc with the same layout.
49            unsafe { System.dealloc(ptr, layout) }
50        }
51    }
52
53    /// Start counting host allocations.
54    pub fn enable() {
55        TRACKING.store(true, Ordering::Release);
56    }
57
58    /// Stop counting host allocations.
59    pub fn disable() {
60        TRACKING.store(false, Ordering::Release);
61    }
62
63    /// Reset the allocation counter to zero.
64    pub fn reset() {
65        ALLOC_COUNT.store(0, Ordering::Release);
66    }
67
68    /// Read the current allocation count.
69    pub fn count() -> usize {
70        ALLOC_COUNT.load(Ordering::Acquire)
71    }
72}
73
74// Stub functions when feature is disabled — optimized away entirely.
75#[cfg(not(feature = "debug-alloc"))]
76pub fn enable() {}
77#[cfg(not(feature = "debug-alloc"))]
78pub fn disable() {}
79#[cfg(not(feature = "debug-alloc"))]
80pub fn reset() {}
81#[cfg(not(feature = "debug-alloc"))]
82pub fn count() -> usize {
83    0
84}
rave_core/debug_alloc.rs

rave_core/
debug_alloc.rs
