memscope_rs/core/
allocator.rs

1//! Custom global allocator for tracking memory allocations.
2
3use std::alloc::{GlobalAlloc, Layout, System};
4
5/// A custom allocator that tracks memory allocations and deallocations.
6///
7/// This allocator wraps the system allocator and records all allocation
8/// and deallocation events through the global memory tracker.
9pub struct TrackingAllocator;
10
11impl TrackingAllocator {
12    /// Create a new tracking allocator instance.
13    pub const fn new() -> Self {
14        Self
15    }
16}
17
18// Thread-local flag to prevent recursive tracking
19thread_local! {
20    static TRACKING_DISABLED: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
21}
22
23unsafe impl GlobalAlloc for TrackingAllocator {
24    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
25        // Allocate memory first
26        let ptr = System.alloc(layout);
27
28        // Track the allocation if it succeeded and tracking is not disabled
29        if !ptr.is_null() {
30            // Check if tracking is disabled for this thread to prevent recursion
31            let should_track = TRACKING_DISABLED.with(|disabled| !disabled.get());
32
33            if should_track {
34                // Temporarily disable tracking to prevent recursion during tracking operations
35                TRACKING_DISABLED.with(|disabled| disabled.set(true));
36
37                // Track the allocation - use try_lock approach to avoid deadlocks
38                if let Ok(tracker) =
39                    std::panic::catch_unwind(crate::core::tracker::get_global_tracker)
40                {
41                    // Ignore errors to prevent allocation failures from breaking the program
42                    let _ = tracker.track_allocation(ptr as usize, layout.size());
43                }
44
45                // Re-enable tracking
46                TRACKING_DISABLED.with(|disabled| disabled.set(false));
47            }
48        }
49
50        ptr
51    }
52
53    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
54        // Track the deallocation first
55        let should_track = TRACKING_DISABLED.with(|disabled| !disabled.get());
56
57        if should_track {
58            // Temporarily disable tracking to prevent recursion
59            TRACKING_DISABLED.with(|disabled| disabled.set(true));
60
61            // Track the deallocation - use try_lock approach to avoid deadlocks
62            if let Ok(tracker) = std::panic::catch_unwind(crate::core::tracker::get_global_tracker)
63            {
64                // Ignore errors to prevent deallocation failures from breaking the program
65                let _ = tracker.track_deallocation(ptr as usize);
66            }
67
68            // Re-enable tracking
69            TRACKING_DISABLED.with(|disabled| disabled.set(false));
70        }
71
72        // Deallocate the memory
73        System.dealloc(ptr, layout);
74    }
75}
76
77impl Default for TrackingAllocator {
78    fn default() -> Self {
79        Self::new()
80    }
81}
memscope_rs/core/allocator.rs

memscope_rs/core/
allocator.rs
