Skip to main content

laminar_core/alloc/
detector.rs

1//! Hot path allocation detector.
2//!
3//! Provides a custom global allocator that tracks allocations and can panic
4//! when allocations occur in marked hot path sections.
5
6#[cfg(feature = "allocation-tracking")]
7use std::alloc::{GlobalAlloc, Layout, System};
8#[cfg(feature = "allocation-tracking")]
9use std::cell::Cell;
10use std::sync::atomic::{AtomicU64, Ordering};
11
12/// Statistics about allocations during hot path detection.
13#[derive(Debug, Default, Clone, Copy)]
14pub struct AllocationStats {
15    /// Total allocations detected in hot path
16    pub hot_path_allocations: u64,
17    /// Total bytes allocated in hot path
18    pub hot_path_bytes: u64,
19    /// Total allocations outside hot path
20    pub normal_allocations: u64,
21    /// Total bytes allocated outside hot path
22    pub normal_bytes: u64,
23}
24
25// Global counters for allocation statistics
26static HOT_PATH_ALLOC_COUNT: AtomicU64 = AtomicU64::new(0);
27static HOT_PATH_ALLOC_BYTES: AtomicU64 = AtomicU64::new(0);
28static NORMAL_ALLOC_COUNT: AtomicU64 = AtomicU64::new(0);
29static NORMAL_ALLOC_BYTES: AtomicU64 = AtomicU64::new(0);
30
31impl AllocationStats {
32    /// Get current allocation statistics.
33    #[must_use]
34    pub fn current() -> Self {
35        Self {
36            hot_path_allocations: HOT_PATH_ALLOC_COUNT.load(Ordering::Relaxed),
37            hot_path_bytes: HOT_PATH_ALLOC_BYTES.load(Ordering::Relaxed),
38            normal_allocations: NORMAL_ALLOC_COUNT.load(Ordering::Relaxed),
39            normal_bytes: NORMAL_ALLOC_BYTES.load(Ordering::Relaxed),
40        }
41    }
42
43    /// Reset all counters to zero.
44    pub fn reset() {
45        HOT_PATH_ALLOC_COUNT.store(0, Ordering::Relaxed);
46        HOT_PATH_ALLOC_BYTES.store(0, Ordering::Relaxed);
47        NORMAL_ALLOC_COUNT.store(0, Ordering::Relaxed);
48        NORMAL_ALLOC_BYTES.store(0, Ordering::Relaxed);
49    }
50}
51
52// Feature-gated hot path detection
53
54#[cfg(feature = "allocation-tracking")]
55thread_local! {
56    /// Nesting level for hot path detection (0 = not in hot path).
57    static HOT_PATH_DEPTH: Cell<usize> = const { Cell::new(0) };
58
59    /// Name of the current hot path section for error messages.
60    static HOT_PATH_SECTION: Cell<Option<&'static str>> = const { Cell::new(None) };
61
62    /// Whether to panic on allocation (vs just count).
63    static PANIC_ON_ALLOC: Cell<bool> = const { Cell::new(true) };
64}
65
66/// Enable hot path detection for current thread.
67///
68/// Supports nesting - each call increments the depth counter.
69///
70/// # Safety
71///
72/// This should only be called through `HotPathGuard`.
73#[cfg(feature = "allocation-tracking")]
74pub(crate) fn enable_hot_path(section: &'static str) {
75    HOT_PATH_DEPTH.with(|d| {
76        let depth = d.get();
77        d.set(depth + 1);
78        // Only update section name on first entry
79        if depth == 0 {
80            HOT_PATH_SECTION.with(|s| s.set(Some(section)));
81        }
82    });
83}
84
85/// Disable hot path detection for current thread.
86///
87/// Supports nesting - only disables when all guards have been dropped.
88#[cfg(feature = "allocation-tracking")]
89pub(crate) fn disable_hot_path() {
90    HOT_PATH_DEPTH.with(|d| {
91        let depth = d.get();
92        if depth > 0 {
93            d.set(depth - 1);
94            if depth == 1 {
95                // Last guard dropped, clear section
96                HOT_PATH_SECTION.with(|s| s.set(None));
97            }
98        }
99    });
100}
101
102/// Check if hot path is currently enabled.
103#[cfg(feature = "allocation-tracking")]
104#[must_use]
105pub fn is_hot_path_enabled() -> bool {
106    HOT_PATH_DEPTH.with(|d| d.get() > 0)
107}
108
109/// Set whether to panic on hot path allocation.
110///
111/// When false, allocations are counted but don't panic.
112/// Useful for collecting stats without crashing.
113#[cfg(feature = "allocation-tracking")]
114pub fn set_panic_on_alloc(panic: bool) {
115    PANIC_ON_ALLOC.with(|p| p.set(panic));
116}
117
118/// Check allocation in hot path and optionally panic.
119#[cfg(feature = "allocation-tracking")]
120#[cold]
121#[inline(never)]
122fn check_hot_path_allocation(op: &str, size: usize) {
123    let is_hot = HOT_PATH_DEPTH.with(|d| d.get() > 0);
124
125    if is_hot {
126        // Record the allocation
127        HOT_PATH_ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
128        HOT_PATH_ALLOC_BYTES.fetch_add(size as u64, Ordering::Relaxed);
129
130        // Check if we should panic
131        let should_panic = PANIC_ON_ALLOC.with(Cell::get);
132        if should_panic {
133            let section = HOT_PATH_SECTION.with(|s| s.get().unwrap_or("unknown"));
134            panic!(
135                "\n\
136                 ╔══════════════════════════════════════════════════════════════╗\n\
137                 ║              ALLOCATION IN HOT PATH DETECTED!                ║\n\
138                 ╠══════════════════════════════════════════════════════════════╣\n\
139                 ║ Operation: {op:50} ║\n\
140                 ║ Size:      {size:50} ║\n\
141                 ║ Section:   {section:50} ║\n\
142                 ╠══════════════════════════════════════════════════════════════╣\n\
143                 ║ Hot path code must be zero-allocation.                       ║\n\
144                 ║ Use pre-allocated buffers, ArrayVec, or ObjectPool instead.  ║\n\
145                 ╚══════════════════════════════════════════════════════════════╝\n",
146                op = op,
147                size = format!("{size} bytes"),
148                section = section
149            );
150        }
151    } else {
152        NORMAL_ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
153        NORMAL_ALLOC_BYTES.fetch_add(size as u64, Ordering::Relaxed);
154    }
155}
156
157// Global allocator (only with feature enabled)
158
159/// Global allocator that detects hot path allocations.
160///
161/// When the `allocation-tracking` feature is enabled and hot path detection
162/// is active for a thread, this allocator will panic on any heap allocation.
163///
164/// # Example
165///
166/// To enable this allocator in your binary:
167///
168/// ```rust,ignore
169/// use laminar_core::alloc::HotPathDetectingAlloc;
170///
171/// #[global_allocator]
172/// static ALLOC: HotPathDetectingAlloc = HotPathDetectingAlloc::new();
173/// ```
174#[cfg(feature = "allocation-tracking")]
175pub struct HotPathDetectingAlloc {
176    inner: System,
177}
178
179#[cfg(feature = "allocation-tracking")]
180impl HotPathDetectingAlloc {
181    /// Create a new hot path detecting allocator.
182    #[must_use]
183    pub const fn new() -> Self {
184        Self { inner: System }
185    }
186}
187
188#[cfg(feature = "allocation-tracking")]
189impl Default for HotPathDetectingAlloc {
190    fn default() -> Self {
191        Self::new()
192    }
193}
194
195#[cfg(feature = "allocation-tracking")]
196// SAFETY: We delegate all allocation operations to the System allocator,
197// which is safe. We only add tracking/panic logic around the calls.
198unsafe impl GlobalAlloc for HotPathDetectingAlloc {
199    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
200        check_hot_path_allocation("alloc", layout.size());
201        // SAFETY: Delegating to System allocator which is safe
202        self.inner.alloc(layout)
203    }
204
205    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
206        // Dealloc is allowed in hot path - we're freeing memory, not allocating
207        // SAFETY: Delegating to System allocator which is safe
208        self.inner.dealloc(ptr, layout);
209    }
210
211    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
212        // Realloc may allocate more memory
213        if new_size > layout.size() {
214            check_hot_path_allocation("realloc", new_size);
215        }
216        // SAFETY: Delegating to System allocator which is safe
217        self.inner.realloc(ptr, layout, new_size)
218    }
219
220    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
221        check_hot_path_allocation("alloc_zeroed", layout.size());
222        // SAFETY: Delegating to System allocator which is safe
223        self.inner.alloc_zeroed(layout)
224    }
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn test_stats_current_and_reset() {
233        AllocationStats::reset();
234
235        let stats = AllocationStats::current();
236        assert_eq!(stats.hot_path_allocations, 0);
237        assert_eq!(stats.normal_allocations, 0);
238    }
239
240    #[test]
241    #[cfg(feature = "allocation-tracking")]
242    fn test_hot_path_enable_disable() {
243        assert!(!is_hot_path_enabled());
244
245        enable_hot_path("test_section");
246        assert!(is_hot_path_enabled());
247
248        disable_hot_path();
249        assert!(!is_hot_path_enabled());
250    }
251}