1use std::alloc::{GlobalAlloc, Layout, System};
2use std::cell::Cell;
3use std::fmt;
4use std::sync::atomic::{AtomicBool, Ordering};
5
6use log::trace;
7
8pub 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
41pub 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}