Poorman's alloc-instrumentation

lib.rs

struct MyAllocator;

#[global_allocator]
static A: MyAllocator = MyAllocator;

const _: () = {
    use ::std::{
        alloc::{
            GlobalAlloc,
            Layout,
            System,
        },
        cell::Cell,
        ops::Not as _,
        sync::{
            Arc,
            atomic::{
                self,
                AtomicUsize,
            },
        },
    };

    use ::dashmap::DashMap;

    type Trace = Arc<[(String, Option<String>, Option<u32>)]>;

    ::lazy_static::lazy_static! {
        static ref TRACKED_ALLOC_SIZES: DashMap<Trace, usize> = DashMap::new();
        static ref TRACKED_ALLOC_PTRS: DashMap<usize, Trace> = DashMap::new();
    }

    static TOTAL_SIZE: AtomicUsize = AtomicUsize::new(0);

    #[track_caller]
    fn get_trace() -> Trace {
        let mut ret = Vec::new();
        ::backtrace::trace(|frame| {
            // let ip = frame.ip();
            // let symbol_address = frame.symbol_address();

            // Resolve this instruction pointer to a symbol name
            ::backtrace::resolve_frame(frame, |symbol| {
                ret.push((
                    symbol
                        .name()
                        .as_ref()
                        .map_or_else(
                            || "<unresolved>".into(),
                            |it| it.to_string(),
                        )
                    ,
                    symbol.filename().map(|it| it.display().to_string()),
                    symbol.lineno(),
                ));
            });

            // keep going to the next frame while:
            ret.len() < 15
        });
        ret.into_boxed_slice().into()
    }

    fn not_reentrantly (f: impl FnOnce())
    {
        thread_local! {
            static REENTRANT: Cell<bool> = false.into();
        }
        if REENTRANT.with(Cell::get).not() {
            REENTRANT.with(|it| it.set(true));
            let () = f();
            REENTRANT.with(|it| it.set(false));
        }
    }

    static MAX_MEMORY: AtomicUsize = AtomicUsize::new(usize::MAX);

    #[ffi_export]
    fn my_alloc_set_limit (bytes: usize)
    {
        MAX_MEMORY.store(dbg!(bytes), atomic::Ordering::Release);
    }

    #[ffi_export]
    fn my_alloc_get_allocated ()
      -> usize
    {
        TOTAL_SIZE.load(atomic::Ordering::Acquire)
    }

    // #[ffi_export]
    fn my_alloc_dump_traced_alloc ()
    {
        (|| -> ::std::io::Result<()> {
            use ::std::io::Write;

            let mut entries = TRACKED_ALLOC_SIZES.iter().map(|guard| {
                let (trace, &total_alloc_size) = guard.pair();
                (trace.clone(), total_alloc_size)
            }).collect::<Vec<_>>();
            entries.sort_unstable_by_key(|&(_, sz)| sz);
            let mut file = ::std::fs::File::create("allocs.trace")?;
            for (trace, mut total_alloc_size) in entries.into_iter().rev() {
                let size = {
                    let total_alloc_mb = total_alloc_size / (1024 * 1024);
                    if total_alloc_mb > 0 {
                        format!("{}.… MiB", total_alloc_mb)
                    } else {
                        let total_alloc_kb = total_alloc_size / 1024;
                        if total_alloc_kb > 0 {
                            format!("{}.… KiB", total_alloc_kb)
                        } else {
                            format!("{} bytes", total_alloc_size)
                        }
                    }
                };
                let () = write!(file, "\n{} from:\n", size)?;
                for caller in &trace[..] {
                    let (fname, mb_file, mb_line) = &*caller;
                    let () = write!(
                        file,
                        concat!(
                            "  {fname}\n",
                            "    at {file}:{line}\n",
                        ),
                        fname = fname,
                        file = mb_file.as_deref().unwrap_or("<unknown>"),
                        line = mb_line.unwrap_or(0),
                    )?;
                }
            }
            let () = file.flush()?;
            let () = file.sync_data()?;
            Ok(())
        })().unwrap();
    }

    unsafe
    impl GlobalAlloc for MyAllocator {
        unsafe
        fn alloc (self: &'_ MyAllocator, layout: Layout)
          -> *mut u8
        {
            let size = layout.size();
            not_reentrantly(|| while {
                let prev = TOTAL_SIZE.load(atomic::Ordering::Relaxed);
                let new = prev.saturating_add(size);
                if new > MAX_MEMORY.load(atomic::Ordering::Acquire) {
                    MAX_MEMORY.store(usize::MAX, atomic::Ordering::Relaxed);
                    my_alloc_dump_traced_alloc();
                    panic!("exhausted allowed memory!");
                }
                TOTAL_SIZE
                    .compare_exchange_weak(
                        prev,
                        new,
                        atomic::Ordering::AcqRel,
                        atomic::Ordering::Relaxed,
                    )
                    .is_err()
            } {});
            let ptr = System.alloc(layout);
            not_reentrantly(|| if ptr.is_null() {
                TOTAL_SIZE.fetch_sub(size, atomic::Ordering::Release);
            } else {
                let trace = get_trace();
                *TRACKED_ALLOC_SIZES.entry(trace.clone()).or_insert(0)
                    += size
                ;
                TRACKED_ALLOC_PTRS.insert(ptr as usize, trace);
            });
            ptr
        }

        unsafe
        fn dealloc (self: &'_ MyAllocator, ptr: *mut u8, layout: Layout)
        {
            not_reentrantly(|| if ptr.is_null().not() {
                let size = layout.size();
                TOTAL_SIZE.fetch_sub(size, atomic::Ordering::Release);
                let (_, trace) = TRACKED_ALLOC_PTRS.remove(&(ptr as usize)).unwrap();
                *TRACKED_ALLOC_SIZES.get_mut(&trace).unwrap()
                    -= size
                ;
            });
            System.dealloc(ptr, layout)
        }
    }
};