Skip to main content

elfo_telemeter/
allocator.rs

1use std::alloc::{GlobalAlloc, Layout};
2
3/// Global allocator providing metrics on allocated memory
4///
5/// ```
6/// # use elfo_telemeter::AllocatorStats;
7/// #[global_allocator]
8/// static ALLOCATOR: AllocatorStats<std::alloc::System> = AllocatorStats::new(std::alloc::System);
9/// ```
10///
11/// Setting this as the global allocator provides two counters:
12/// `elfo_allocated_bytes_total` and `elfo_deallocated_bytes_total`, tracking
13/// total allocated and deallocated memory in bytes.
14#[instability::unstable]
15pub struct AllocatorStats<A> {
16    inner: A,
17}
18
19impl<A> AllocatorStats<A> {
20    /// Wrap a global allocator, instrumenting it with metrics
21    pub const fn new(inner: A) -> Self {
22        Self { inner }
23    }
24}
25
26// SAFETY: it augmentes the logic of an inner allocator, but does not change it.
27unsafe impl<A> GlobalAlloc for AllocatorStats<A>
28where
29    A: GlobalAlloc,
30{
31    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
32        let ptr = self.inner.alloc(layout);
33        if !ptr.is_null() {
34            elfo_core::scope::try_with(|scope| {
35                // TODO: a separate counter for messaging.
36                scope.increment_allocated_bytes(layout.size());
37            });
38        }
39        ptr
40    }
41
42    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
43        self.inner.dealloc(ptr, layout);
44        elfo_core::scope::try_with(|scope| {
45            // TODO: a separate counter for messaging.
46            scope.increment_deallocated_bytes(layout.size());
47        });
48    }
49
50    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
51        let ptr = self.inner.alloc_zeroed(layout);
52        if !ptr.is_null() {
53            elfo_core::scope::try_with(|scope| {
54                scope.increment_allocated_bytes(layout.size());
55            });
56        }
57        ptr
58    }
59
60    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
61        let ptr = self.inner.realloc(ptr, layout, new_size);
62        if !ptr.is_null() {
63            elfo_core::scope::try_with(|scope| {
64                scope.increment_deallocated_bytes(layout.size());
65                scope.increment_allocated_bytes(new_size);
66            });
67        }
68        ptr
69    }
70}