use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicU64, Ordering};
struct CountingAllocator;
static ALLOCATION_CALLS: AtomicU64 = AtomicU64::new(0);
static DEALLOCATION_CALLS: AtomicU64 = AtomicU64::new(0);
static REQUESTED_BYTES: AtomicU64 = AtomicU64::new(0);
unsafe impl GlobalAlloc for CountingAllocator {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
ALLOCATION_CALLS.fetch_add(1, Ordering::Relaxed);
REQUESTED_BYTES.fetch_add(size_u64(layout.size()), Ordering::Relaxed);
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
DEALLOCATION_CALLS.fetch_add(1, Ordering::Relaxed);
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
ALLOCATION_CALLS.fetch_add(1, Ordering::Relaxed);
REQUESTED_BYTES.fetch_add(size_u64(new_size), Ordering::Relaxed);
unsafe { System.realloc(ptr, layout, new_size) }
}
}
#[global_allocator]
static GLOBAL: CountingAllocator = CountingAllocator;
#[derive(Clone, Copy)]
pub(crate) struct AllocationSnapshot {
allocation_calls: u64,
deallocation_calls: u64,
requested_bytes: u64,
}
#[derive(Clone, Copy, Default, serde::Serialize)]
pub(crate) struct AllocationProfile {
pub(crate) allocation_calls: u64,
pub(crate) deallocation_calls: u64,
pub(crate) requested_bytes: u64,
}
impl AllocationSnapshot {
pub(crate) fn now() -> Self {
Self {
allocation_calls: ALLOCATION_CALLS.load(Ordering::Relaxed),
deallocation_calls: DEALLOCATION_CALLS.load(Ordering::Relaxed),
requested_bytes: REQUESTED_BYTES.load(Ordering::Relaxed),
}
}
pub(crate) fn elapsed(self) -> AllocationProfile {
let end = Self::now();
AllocationProfile {
allocation_calls: end.allocation_calls.wrapping_sub(self.allocation_calls),
deallocation_calls: end.deallocation_calls.wrapping_sub(self.deallocation_calls),
requested_bytes: end.requested_bytes.wrapping_sub(self.requested_bytes),
}
}
}
fn size_u64(size: usize) -> u64 {
u64::try_from(size).unwrap_or(u64::MAX)
}