memscope_rs/
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) = std::panic::catch_unwind(crate::tracker::get_global_tracker) {
39                    // Ignore errors to prevent allocation failures from breaking the program
40                    let _ = tracker.track_allocation(ptr as usize, layout.size());
41                }
42
43                // Re-enable tracking
44                TRACKING_DISABLED.with(|disabled| disabled.set(false));
45            }
46        }
47
48        ptr
49    }
50
51    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
52        // Track the deallocation first
53        let should_track = TRACKING_DISABLED.with(|disabled| !disabled.get());
54
55        if should_track {
56            // Temporarily disable tracking to prevent recursion
57            TRACKING_DISABLED.with(|disabled| disabled.set(true));
58
59            // Track the deallocation - use try_lock approach to avoid deadlocks
60            if let Ok(tracker) = std::panic::catch_unwind(crate::tracker::get_global_tracker) {
61                // Ignore errors to prevent deallocation failures from breaking the program
62                let _ = tracker.track_deallocation(ptr as usize);
63            }
64
65            // Re-enable tracking
66            TRACKING_DISABLED.with(|disabled| disabled.set(false));
67        }
68
69        // Deallocate the memory
70        System.dealloc(ptr, layout);
71    }
72}
73
74impl Default for TrackingAllocator {
75    fn default() -> Self {
76        Self::new()
77    }
78}
memscope_rs/allocator.rs

memscope_rs/
allocator.rs
