nordisk/nornir

crates/mimir-ingest/src/alloc_count.rs

//! Test builds only: the global allocator counts, per thread, the bytes it
//! has handed out and not had back — what a structure built on the test's
//! own thread holds, whatever other tests do meanwhile on theirs.

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;

thread_local! {
    static NET: Cell<isize> = const { Cell::new(0) };
}

fn add(n: isize) {
    // a thread being torn down has no counter left: nothing to count for it
    let _ = NET.try_with(|c| c.set(c.get() + n));
}

struct Counting;

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(l) };
        if !p.is_null() {
            add(l.size() as isize);
        }
        p
    }
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc_zeroed(l) };
        if !p.is_null() {
            add(l.size() as isize);
        }
        p
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) };
        add(-(l.size() as isize));
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
        let q = unsafe { System.realloc(p, l, new) };
        if !q.is_null() {
            add(new as isize - l.size() as isize);
        }
        q
    }
}

#[global_allocator]
static COUNTING: Counting = Counting;

/// Bytes this thread holds from the allocator (allocated minus freed).
pub(crate) fn held() -> isize {
    NET.with(|c| c.get())
}