Skip to content

Instantly share code, notes, and snippets.

@pftbest
Created January 9, 2024 00:10
Show Gist options
  • Save pftbest/d8e997ce2fb045c166135e6045f46906 to your computer and use it in GitHub Desktop.
Save pftbest/d8e997ce2fb045c166135e6045f46906 to your computer and use it in GitHub Desktop.
backtracing allocator
use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::RefCell;
use std::fs::File;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Mutex;
use std::thread_local;
thread_local! {
static BT_ALLOC_BUSY: RefCell<bool> = RefCell::new(false);
}
struct BacktracingAllocatorData {
file: Option<File>,
}
pub struct BacktracingAllocator {
active: AtomicBool,
data: Mutex<BacktracingAllocatorData>,
}
impl BacktracingAllocator {
pub const fn new() -> Self {
Self {
active: AtomicBool::new(false),
data: Mutex::new(BacktracingAllocatorData { file: None }),
}
}
pub fn initialize(&self, f: File) {
let mut data = self.data.lock().unwrap();
data.file.replace(f);
}
pub fn activate(&self) {
self.active.store(true, Ordering::Relaxed);
}
pub fn deactivate(&self) {
self.active.store(false, Ordering::Relaxed);
}
fn print_size(&self, address: usize, size: usize, action: Action) {
let active = self.active.load(Ordering::Relaxed);
let unwinding = BT_ALLOC_BUSY.with_borrow(|f| *f);
if !active || unwinding {
return;
}
let mut data = self.data.lock().unwrap();
BT_ALLOC_BUSY.with_borrow_mut(|b| *b = true);
if let Some(file) = data.file.as_mut() {
let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as usize;
let act = match action {
Action::Allocating => 'A',
Action::Deallocating => 'D',
};
let _ = writeln!(file, "\n{:016X} {} {:016X} {:016X}", timestamp, act, address, size);
if action == Action::Allocating {
let bt = backtrace::Backtrace::new();
let mut trace = String::new();
for frame in bt.frames().iter().skip(3) {
for sym in frame.symbols() {
let file = sym.filename().map(|n| n.as_os_str().as_bytes()).unwrap_or(b"?");
let file = String::from_utf8_lossy(file);
let line = sym.lineno().unwrap_or(0);
let _ = writeln!(&mut trace, "=> {:?}\n {}:{}", sym.name(), file, line);
}
}
let _ = writeln!(file, "{}", trace);
}
}
BT_ALLOC_BUSY.with_borrow_mut(|b| *b = false);
}
}
#[derive(Eq, PartialEq)]
enum Action {
Allocating,
Deallocating,
}
unsafe impl GlobalAlloc for BacktracingAllocator {
#[track_caller]
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let ptr = System.alloc(layout);
self.print_size(ptr as usize, layout.size(), Action::Allocating);
ptr
}
#[track_caller]
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.print_size(ptr as usize, layout.size(), Action::Deallocating);
System.dealloc(ptr, layout);
}
}
#[global_allocator]
static BT_ALLOC: BacktracingAllocator = BacktracingAllocator::new();
fn main() -> Result<(), Box<dyn Error>> {
let f = File::create("trace.txt").unwrap();
BT_ALLOC.initialize(f);
BT_ALLOC.activate();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment