tracing_allocations/
lib.rs

1use core::mem;
2use core::ops::DerefMut;
3
4use core::{
5    alloc::{GlobalAlloc, Layout},
6    cell::RefCell,
7};
8
9thread_local! {
10    /// Flag controlling whether to emit tracing events for allocations/deallocations on this thread.
11    pub static TRACE_ALLOCATOR: RefCell<bool> = RefCell::new(false);
12}
13
14/// An allocator that emits tracing events.
15///
16/// ## Usage
17/// ```
18/// use std::alloc::System;
19/// use tracing_allocations::TracingAllocator;
20///
21/// #[global_allocator]
22/// static ALLOCATOR: TracingAllocator<System> = TracingAllocator::new(System);
23/// ```
24#[non_exhaustive]
25pub struct TracingAllocator<A> {
26    pub allocator: A,
27}
28
29impl<A> TracingAllocator<A> {
30    /// Constructs a tracing allocator.
31    ///
32    /// ## Usage
33    /// ```
34    /// use std::alloc::System;
35    /// use tracing_allocations::TracingAllocator;
36    ///
37    /// #[global_allocator]
38    /// static ALLOCATOR: TracingAllocator<System> = TracingAllocator::new(System);
39    /// ```
40    pub const fn new(allocator: A) -> Self {
41        Self { allocator }
42    }
43}
44
45unsafe impl<A> GlobalAlloc for TracingAllocator<A>
46where
47    A: GlobalAlloc,
48{
49    #[track_caller]
50    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
51        let ptr = self.allocator.alloc(layout);
52        // trace the allocation
53        let _ = TRACE_ALLOCATOR.try_with(|guard| {
54            // `guard.try_borrow_mut()` prevents us from tracing our traces
55            if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
56                if mem::replace(trace_allocations.deref_mut(), false) {
57                    tracing::trace! {
58                        addr = ptr as usize,
59                        size = layout.size(),
60                        "alloc",
61                    };
62                    *trace_allocations = true;
63                }
64                drop(trace_allocations);
65            }
66        });
67        ptr
68    }
69
70    #[track_caller]
71    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
72        self.allocator.dealloc(ptr, layout);
73        let _ = TRACE_ALLOCATOR.try_with(|guard| {
74            // `guard.try_borrow_mut()` prevents us from tracing our traces
75            if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
76                if mem::replace(trace_allocations.deref_mut(), false) {
77                    tracing::trace! {
78                        addr = ptr as usize,
79                        size = layout.size(),
80                        "dealloc",
81                    };
82                    *trace_allocations = true;
83                }
84                drop(guard);
85            }
86        });
87    }
88}
89
90/// Trace allocations occurring within `f`.
91pub fn trace_allocations<F: FnOnce() -> R, R>(f: F) -> R {
92    TRACE_ALLOCATOR.with(|guard| {
93        let mut previous_state = false;
94        if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
95            previous_state = mem::replace(&mut trace_allocations, true);
96        }
97        let res = f();
98        if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
99            *trace_allocations = previous_state;
100        }
101        res
102    })
103}
104
105/// Do not trace allocations occurring within `f`.
106pub fn ignore_allocations<F: FnOnce() -> R, R>(f: F) -> R {
107    TRACE_ALLOCATOR.with(|guard| {
108        let mut previous_state = true;
109        if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
110            previous_state = mem::replace(&mut trace_allocations, false);
111        }
112        let res = f();
113        if let Ok(mut trace_allocations) = guard.try_borrow_mut() {
114            *trace_allocations = previous_state;
115        }
116        res
117    })
118}
tracing_allocations/lib.rs

tracing_allocations/
lib.rs
