48 lines
1.2 KiB
Rust
48 lines
1.2 KiB
Rust
//! A small global allocator implementation which is intended to keep track of
|
|
//! the number of allocated bytes to ensure that all our integration glue indeed
|
|
//! manages memory correctly and doesn't leak anything.
|
|
|
|
use std::alloc::{GlobalAlloc, Layout, System};
|
|
use std::sync::atomic::{AtomicUsize, Ordering::SeqCst};
|
|
|
|
#[global_allocator]
|
|
static ALLOC: A = A;
|
|
|
|
static ALLOC_AMT: AtomicUsize = AtomicUsize::new(0);
|
|
|
|
struct A;
|
|
|
|
unsafe impl GlobalAlloc for A {
|
|
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
|
let ptr = System.alloc(layout);
|
|
if !ptr.is_null() {
|
|
ALLOC_AMT.fetch_add(layout.size(), SeqCst);
|
|
}
|
|
return ptr;
|
|
}
|
|
|
|
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
|
// Poison all deallocations to try to catch any use-after-free in the
|
|
// bindings as early as possible.
|
|
std::ptr::write_bytes(ptr, 0xde, layout.size());
|
|
ALLOC_AMT.fetch_sub(layout.size(), SeqCst);
|
|
System.dealloc(ptr, layout)
|
|
}
|
|
}
|
|
|
|
pub fn get() -> usize {
|
|
ALLOC_AMT.load(SeqCst)
|
|
}
|
|
|
|
pub fn guard() -> impl Drop {
|
|
struct A(usize);
|
|
|
|
impl Drop for A {
|
|
fn drop(&mut self) {
|
|
assert_eq!(get(), self.0);
|
|
}
|
|
}
|
|
|
|
A(get())
|
|
}
|