readings_probe/
alloc.rs

1use std::sync::atomic::AtomicUsize;
2
3pub static ALLOCATED: AtomicUsize = AtomicUsize::new(0);
4pub static FREEED: AtomicUsize = AtomicUsize::new(0);
5
6/// Setup global allocator instrumentation, to track rust-managed memory.
7///
8/// It is not mandatory to do so, as we also register the RSZ and VSZ as
9/// reported by the OS, but it may be interesting. From our experience it may be
10/// worth activating it as the cost is relatively modest.
11///
12/// This macro allows to specify a specific allocator instance to wrap and make global
13/// (for instance, jemalloc).
14#[macro_export]
15macro_rules! wrap_global_allocator {
16    ($alloc:path) => {
17        #[global_allocator]
18        static A: InstrumentedAllocator = InstrumentedAllocator;
19
20        struct InstrumentedAllocator;
21
22        unsafe impl std::alloc::GlobalAlloc for InstrumentedAllocator {
23            unsafe fn alloc(&self, layout: std::alloc::Layout) -> *mut u8 {
24                use std::sync::atomic::Ordering::Relaxed;
25                let ptr = $alloc.alloc(layout);
26                if !ptr.is_null() {
27                    $crate::alloc::ALLOCATED.fetch_add(layout.size(), Relaxed);
28                }
29                ptr
30            }
31            unsafe fn dealloc(&self, ptr: *mut u8, layout: std::alloc::Layout) {
32                use std::sync::atomic::Ordering::Relaxed;
33                if !ptr.is_null() {
34                    $crate::alloc::FREEED.fetch_add(layout.size(), Relaxed);
35                }
36                $alloc.dealloc(ptr, layout);
37            }
38            unsafe fn realloc(&self, ptr: *mut u8, layout: std::alloc::Layout, new_size: usize) -> *mut u8 {
39                use std::sync::atomic::Ordering::*;
40                if !ptr.is_null() {
41                    $crate::alloc::FREEED.fetch_add(layout.size(), Relaxed);
42                }
43                let ptr = $alloc.realloc(ptr, layout, new_size);
44                if !ptr.is_null() {
45                    $crate::alloc::ALLOCATED.fetch_add(new_size, Relaxed);
46                }
47                ptr
48            }
49        }
50    };
51}
52
53/// Setup global allocator instrumentation, to track rust-managed memory.
54///
55/// It is not mandatory to do so, as we also register the RSZ and VSZ as
56/// reported by the OS, but it may be interesting. From our experience it may be
57/// worth activating it as the cost is relatively modest.
58///
59/// This macro wrap the System allocator. Use `wrap_global_allocator` to wrap
60/// another global allocator.
61#[macro_export]
62macro_rules! instrumented_allocator {
63    () => {
64        $crate::wrap_global_allocator!(std::alloc::System);
65    };
66}
readings_probe/alloc.rs

readings_probe/
alloc.rs
