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
use std::alloc::{GlobalAlloc, Layout, System}; use std::sync::atomic::{AtomicUsize, Ordering::SeqCst}; #[derive(Debug)] pub struct CountingAllocator { counter: AtomicUsize, } impl CountingAllocator { pub const fn new() -> Self { Self { counter: AtomicUsize::new(0), } } pub fn current_usage(&self) -> usize { self.counter.load(SeqCst) } } unsafe impl GlobalAlloc for CountingAllocator { unsafe fn alloc(&self, layout: Layout) -> *mut u8 { let ret = System.alloc(layout); if !ret.is_null() { self.counter.fetch_add(layout.size(), SeqCst); } return ret; } unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { System.dealloc(ptr, layout); self.counter.fetch_sub(layout.size(), SeqCst); } }