rs_malloc_tracker/
lib.rs

1#![deny(warnings)]
2
3use std::{
4    alloc::{GlobalAlloc, Layout},
5    ffi::{CStr, c_int, c_void},
6    sync::atomic::AtomicUsize,
7};
8
9use crate::consumer::{ZoneKind, remove_zone, send_zone};
10
11mod consumer;
12mod macros;
13#[cfg(feature = "mmap")]
14mod mmap;
15mod prometheus;
16mod thread_local_sender;
17
18#[cfg(test)]
19mod tests;
20
21#[macro_export]
22macro_rules! c_stringify {
23    ($x:expr) => {
24        const {
25            let bytes = concat!($x, "\0").as_bytes();
26            match ::core::ffi::CStr::from_bytes_with_nul(bytes) {
27                Ok(cstr) => cstr,
28                Err(_) => unreachable!(),
29            }
30        }
31    };
32}
33
34pub struct AtomicFnPtr<F> {
35    ptr: AtomicUsize,
36    shadow_name: Option<&'static CStr>,
37    _t: std::marker::PhantomData<F>,
38}
39
40fn_impl!(extern "C" fn(usize) -> *mut c_void);
41fn_impl!(extern "C" fn(usize, usize) -> *mut c_void);
42fn_impl!(extern "C" fn(*mut *mut c_void, usize, usize) -> c_int);
43fn_impl!(extern "C" fn(*mut c_void) -> ());
44fn_impl!(extern "C" fn(*mut c_void, usize) -> *mut c_void);
45fn_impl!(extern "C" fn(old_ptr: *mut c_void, number: usize, size: usize) -> *mut c_void);
46
47shadow!(malloc(size: usize) -> *mut c_void {
48    let ptr = f!()(size);
49
50    send_zone(ZoneKind::Malloc, ptr, size);
51
52    ptr
53});
54
55shadow!(calloc(number: usize, size: usize) -> *mut c_void {
56    let ptr = f!()(number, size);
57
58    send_zone(ZoneKind::Calloc, ptr, number * size);
59
60    ptr
61});
62
63shadow!(posix_memalign(ptr: *mut *mut c_void, alignment: usize, size: usize) -> c_int {
64    let ret = f!()(ptr, alignment, size);
65
66    if ret == 0 && let Some(ptr) = unsafe { ptr.as_mut() } {
67        send_zone(ZoneKind::Memalign, unsafe { *ptr }, size);
68    }
69
70    ret
71});
72
73shadow!(memalign(alignment: usize, size: usize) -> *mut c_void {
74    let ptr = f!()(alignment, size);
75
76    send_zone(ZoneKind::Memalign, ptr, size);
77
78    ptr
79});
80
81shadow!(aligned_alloc(alignment: usize, size: usize) -> *mut c_void {
82    let ptr = f!()(alignment, size);
83
84    send_zone(ZoneKind::Memalign, ptr, size);
85
86    ptr
87});
88
89shadow!(pvalloc(size: usize) -> *mut c_void {
90    let ptr = f!()(size);
91
92    send_zone(ZoneKind::Memalign, ptr, size);
93
94    ptr
95});
96
97shadow!(valloc(size: usize) -> *mut c_void {
98    let ptr = f!()(size);
99
100    send_zone(ZoneKind::Memalign, ptr, size);
101
102    ptr
103});
104
105shadow!(realloc(old_ptr: *mut c_void, size: usize) -> *mut c_void {
106    let new_ptr = f!()(old_ptr, size);
107
108    remove_zone(old_ptr);
109    if !new_ptr.is_null() {
110        send_zone(ZoneKind::Realloc, new_ptr, size);
111    }
112
113    new_ptr
114});
115
116shadow!(reallocarray(old_ptr: *mut c_void, number: usize, size: usize) -> *mut c_void {
117    let new_ptr = f!()(old_ptr, number, size);
118
119    remove_zone(old_ptr);
120    if !new_ptr.is_null() {
121        send_zone(ZoneKind::Realloc, new_ptr, number * size);
122    }
123
124    new_ptr
125});
126
127shadow!(free(ptr: *mut c_void) -> () {
128    f!()(ptr);
129
130    remove_zone(ptr);
131});
132
133/*
134 * Ensure we do not use our own shadowed functions within this library.
135 * Recursion would be bad :)
136 *
137 * So we use a separate copy of malloc() and free() for our GlobalAlloc
138 */
139
140static RUST_MALLOC: AtomicFnPtr<extern "C" fn(usize) -> *mut c_void> =
141    AtomicFnPtr::<extern "C" fn(usize) -> *mut c_void>::new(Some(c_stringify!("malloc")));
142static RUST_FREE: AtomicFnPtr<extern "C" fn(*mut c_void)> =
143    AtomicFnPtr::<extern "C" fn(*mut c_void)>::new(Some(c_stringify!("free")));
144
145#[unsafe(no_mangle)]
146pub fn rust_malloc(size: usize) -> *mut u8 {
147    RUST_MALLOC.load_or_shadow()(size) as *mut u8
148}
149
150#[unsafe(no_mangle)]
151pub fn rust_free(ptr: *mut u8) {
152    RUST_FREE.load_or_shadow()(ptr as *mut c_void)
153}
154
155struct MyAllocator;
156unsafe impl GlobalAlloc for MyAllocator {
157    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
158        rust_malloc(layout.size())
159    }
160
161    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
162        rust_free(ptr);
163    }
164}
165#[global_allocator]
166static GLOBAL_ALLOC: MyAllocator = MyAllocator;
rs_malloc_tracker/lib.rs

rs_malloc_tracker/
lib.rs
