1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
/// Prototype of ring-fenced memory allocators for `thag_profiler`.
///
/// The `global_allocator` attribute flags a `Dispatcher` which dispatches each
/// memory allocation, deallocation and reallocation requests to one of two allocators
/// according to the designated current allocator (held in an atomic boolean) at the
/// moment that it receives the request. The default allocator is `TaskAware` and is
/// used for user code, while the regular system allocator `System` handles requests
/// from profiler code. The role of the `TaskAware` allocator is to record the details
/// of the user code allocation events before passing them to the system allocator.
///
/// To invoke the system allocator directly, profiler code must call a function or
/// closure with fn `with_sys_alloc`, which checks the current allocator, and if it
/// finds it to be `TaskAware`, changes it to `System` and runs the function or closure,
/// with a guard to restore the default to `TaskAware`. If the current allocator is
/// already `System`, `with_sys_alloc` concludes that it must be running nested under
/// another `with_sys_alloc` call, so does nothing except run the function or closure.
///
/// The flaw in this design is its vulnerability to race conditions, e.g. user code
/// in another thread could fail to go through `TaskAware` if `with_sys_alloc` is
/// running concurrently, or conversely an outer `with_sys_alloc` ending in one thread
/// could prematurely reset the current allocator to `TaskAware` while another
/// instance is still running in another thread. We can and do build in a check in
/// the TaskAware branch to detect and ignore profiler code, but in practice there is
/// little sign of such races being a problem.
///
/// Attempts to resolve this issue with thread-local storage have not borne fruit.
/// For instance async tasks are by no means guaranteed to resume in the same thread
/// after suspension.
/// The ideal would seem to be a reentrant Mutex or RwLock with mutability - so far tried
/// without success, but a subject for another prototype.
//# Purpose: Prototype of a ring-fenced allocator for memory profiling.
//# Categories: profiling, prototype
use std::{
alloc::{GlobalAlloc, Layout, System},
fmt,
sync::atomic::{AtomicBool, Ordering},
};
static USING_SYSTEM_ALLOCATOR: AtomicBool = AtomicBool::new(false);
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Allocator {
/// Task-aware allocator that tracks which task allocated memory
TaskAware,
/// System allocator for profiling operations
System,
}
impl fmt::Display for Allocator {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Allocator::TaskAware => write!(f, "TaskAware"),
Allocator::System => write!(f, "System"),
}
}
}
pub fn with_sys_alloc<F, R>(f: F) -> R
where
F: FnOnce() -> R,
{
if current_allocator() == Allocator::System {
eprintln!("Already in SystemAllocator");
return f();
}
USING_SYSTEM_ALLOCATOR.store(true, Ordering::SeqCst);
// Create struct to handle cleanup on drop
struct Cleanup;
impl Drop for Cleanup {
fn drop(&mut self) {
USING_SYSTEM_ALLOCATOR.store(false, Ordering::SeqCst);
}
}
// Create guard to restore on scope exit
let _cleanup = Cleanup {};
// Run the function
eprintln!("Switched to SystemAllocator");
f()
}
pub fn current_allocator() -> Allocator {
if USING_SYSTEM_ALLOCATOR.load(Ordering::Relaxed) {
// eprintln!("Using system allocator");
Allocator::System
} else {
Allocator::TaskAware
}
}
// Create a direct static instance
#[global_allocator]
static ALLOCATOR: Dispatcher = Dispatcher::new();
/// Dispatcher that routes allocation requests to the active allocator
/// according to the USING_SYSTEM_ALLOCATOR variable for the current thread.
pub struct Dispatcher {
pub task_aware: TaskAwareAllocator,
pub system: std::alloc::System,
}
impl Dispatcher {
pub const fn new() -> Self {
Self {
task_aware: TaskAwareAllocator,
system: std::alloc::System,
}
}
}
impl Default for Dispatcher {
fn default() -> Self {
Self::new()
}
}
unsafe impl GlobalAlloc for Dispatcher {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let current = current_allocator();
match current {
Allocator::System => unsafe { self.system.alloc(layout) },
Allocator::TaskAware => unsafe { self.task_aware.alloc(layout) },
}
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
if ptr.is_null() {
return;
}
match current_allocator() {
Allocator::System => unsafe { self.system.dealloc(ptr, layout) },
Allocator::TaskAware => unsafe { self.task_aware.dealloc(ptr, layout) },
}
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
if ptr.is_null() {
return unsafe {
self.alloc(Layout::from_size_align_unchecked(new_size, layout.align()))
};
}
match current_allocator() {
Allocator::System => unsafe { self.system.realloc(ptr, layout, new_size) },
Allocator::TaskAware => unsafe { self.task_aware.realloc(ptr, layout, new_size) },
}
}
}
/// Task-aware allocator that tracks memory allocations
pub struct TaskAwareAllocator;
// Static instance for global access
#[allow(dead_code)]
static TASK_AWARE_ALLOCATOR: TaskAwareAllocator = TaskAwareAllocator;
unsafe impl GlobalAlloc for TaskAwareAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
with_sys_alloc(|| {
let ptr = unsafe { System.alloc(layout) };
println!("In TaskAwareAllocator for size {}", layout.size());
ptr
})
}
#[allow(clippy::too_many_lines)]
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
with_sys_alloc(|| {
if !ptr.is_null() {
// println!("In system deallocator for size {}", layout.size());
// Forward to system allocator for deallocation
unsafe { System.dealloc(ptr, layout) };
}
});
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, _new_size: usize) -> *mut u8 {
if !ptr.is_null() {
with_sys_alloc(|| {
println!("In system reallocator for size {}", layout.size());
});
}
ptr
}
}
fn main() {
let data1: Vec<u8> = vec![0; 1024];
println!("1. current_allocator()={}", current_allocator());
let data2: Vec<u8> = with_sys_alloc(|| {
with_sys_alloc(|| println!("Nested sys alloc"));
vec![0; 2048]
});
println!("2. current_allocator()={}", current_allocator());
with_sys_alloc(|| println!("data1.len()={}, data2.len()={}", data1.len(), data2.len()));
println!("3. current_allocator()={}", current_allocator());
}