backtracing allocator

backtracing_alloc.rs

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();
}