logging_allocator/
lib.rs

1use std::alloc::{GlobalAlloc, Layout, System};
2use std::cell::Cell;
3use std::fmt;
4use std::sync::atomic::{AtomicBool, Ordering};
5
6use log::trace;
7
8/// A wrapper allocator that logs messages on allocation.
9pub struct LoggingAllocator<A = System> {
10    enabled: AtomicBool,
11    allocator: A,
12}
13
14impl LoggingAllocator<System> {
15    pub const fn new() -> Self {
16        LoggingAllocator::with_allocator(System)
17    }
18}
19
20impl<A> LoggingAllocator<A> {
21    pub const fn with_allocator(allocator: A) -> Self {
22        LoggingAllocator {
23            enabled: AtomicBool::new(false),
24            allocator,
25        }
26    }
27
28    pub fn enable_logging(&self) {
29        self.enabled.store(true, Ordering::SeqCst)
30    }
31
32    pub fn disable_logging(&self) {
33        self.enabled.store(false, Ordering::SeqCst)
34    }
35
36    pub fn logging_enabled(&self) -> bool {
37        self.enabled.load(Ordering::SeqCst)
38    }
39}
40
41/// Execute a closure without logging on allocations.
42pub fn run_guarded<F>(f: F)
43where
44    F: FnOnce(),
45{
46    thread_local! {
47        static GUARD: Cell<bool> = Cell::new(false);
48    }
49
50    GUARD.with(|guard| {
51        if !guard.replace(true) {
52            f();
53            guard.set(false)
54        }
55    })
56}
57
58unsafe impl<A> GlobalAlloc for LoggingAllocator<A>
59where
60    A: GlobalAlloc,
61{
62    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
63        let ptr = self.allocator.alloc(layout);
64        if self.logging_enabled() {
65            run_guarded(|| {
66                trace!("alloc {}", Fmt(ptr, layout.size(), layout.align()));
67            });
68        }
69        ptr
70    }
71
72    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
73        self.allocator.dealloc(ptr, layout);
74        if self.logging_enabled() {
75            run_guarded(|| trace!("dealloc {}", Fmt(ptr, layout.size(), layout.align()),));
76        }
77    }
78
79    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
80        let ptr = self.allocator.alloc_zeroed(layout);
81        if self.logging_enabled() {
82            run_guarded(|| {
83                trace!("alloc_zeroed {}", Fmt(ptr, layout.size(), layout.align()));
84            });
85        }
86        ptr
87    }
88
89    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
90        let new_ptr = self.allocator.realloc(ptr, layout, new_size);
91        if self.logging_enabled() {
92            run_guarded(|| {
93                trace!(
94                    "realloc {} to {}",
95                    Fmt(ptr, layout.size(), layout.align()),
96                    Fmt(new_ptr, new_size, layout.align())
97                );
98            });
99        }
100        new_ptr
101    }
102}
103
104struct Fmt(*mut u8, usize, usize);
105
106impl fmt::Display for Fmt {
107    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
108        write!(
109            f,
110            "[address={:p}, size={}, align={}]",
111            self.0, self.1, self.2
112        )
113    }
114}
logging_allocator/lib.rs

logging_allocator/
lib.rs
