Skip to content

Instantly share code, notes, and snippets.

@mosmeh
Created September 27, 2021 22:24
Show Gist options
  • Save mosmeh/a502c19d85adb951b6fa2f3a960dafea to your computer and use it in GitHub Desktop.
Save mosmeh/a502c19d85adb951b6fa2f3a960dafea to your computer and use it in GitHub Desktop.
use std::{
alloc::{GlobalAlloc, Layout, System},
sync::atomic::{AtomicUsize, Ordering},
};
pub struct Allocator {
inner: System,
counter: AtomicUsize,
}
unsafe impl GlobalAlloc for Allocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
self.counter.fetch_add(layout.size(), Ordering::SeqCst);
self.inner.alloc(layout)
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
self.inner.dealloc(ptr, layout);
self.counter.fetch_sub(layout.size(), Ordering::SeqCst);
}
}
impl Allocator {
pub const fn new() -> Self {
Self {
inner: System,
counter: AtomicUsize::new(0),
}
}
pub fn reset(&self) {
self.counter.store(0, Ordering::SeqCst);
}
pub fn get(&self) -> usize {
self.counter.load(Ordering::SeqCst)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment