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
35
36

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicUsize, Ordering};

#[macro_export]
macro_rules! track_mem {
    () => {
        #[global_allocator]
        static GLOBAL: heaptrack::TrackingAllocator = heaptrack::TrackingAllocator;
    };
}

static BYTES_REQUESTED_ALIGNED: AtomicUsize = AtomicUsize::new(0);

pub fn bytes_requested_aligned() -> usize {
    BYTES_REQUESTED_ALIGNED.load(Ordering::SeqCst)
}

pub struct TrackingAllocator;

unsafe impl GlobalAlloc for TrackingAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let requested = layout.size();
        let align = layout.align();
        let requested_aligned = (requested as *const u8).align_offset(align) + requested;
        BYTES_REQUESTED_ALIGNED.fetch_add(requested_aligned, Ordering::SeqCst);
        System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let requested = layout.size();
        let align = layout.align();
        let requested_aligned = (requested as *const u8).align_offset(align) + requested;
        BYTES_REQUESTED_ALIGNED.fetch_sub(requested_aligned, Ordering::SeqCst);
        System.dealloc(ptr, layout)
    }
}