rs_malloc_tracker 1.0.1

#![deny(warnings)]

use std::{
    alloc::{GlobalAlloc, Layout},
    ffi::{CStr, c_int, c_void},
    sync::atomic::AtomicUsize,
};

use crate::consumer::{ZoneKind, remove_zone, send_zone};

mod consumer;
mod macros;
#[cfg(feature = "mmap")]
mod mmap;
mod prometheus;
mod thread_local_sender;

#[cfg(test)]
mod tests;

#[macro_export]
macro_rules! c_stringify {
    ($x:expr) => {
        const {
            let bytes = concat!($x, "\0").as_bytes();
            match ::core::ffi::CStr::from_bytes_with_nul(bytes) {
                Ok(cstr) => cstr,
                Err(_) => unreachable!(),
            }
        }
    };
}

pub struct AtomicFnPtr<F> {
    ptr: AtomicUsize,
    shadow_name: Option<&'static CStr>,
    _t: std::marker::PhantomData<F>,
}

fn_impl!(extern "C" fn(usize) -> *mut c_void);
fn_impl!(extern "C" fn(usize, usize) -> *mut c_void);
fn_impl!(extern "C" fn(*mut *mut c_void, usize, usize) -> c_int);
fn_impl!(extern "C" fn(*mut c_void) -> ());
fn_impl!(extern "C" fn(*mut c_void, usize) -> *mut c_void);
fn_impl!(extern "C" fn(old_ptr: *mut c_void, number: usize, size: usize) -> *mut c_void);

shadow!(malloc(size: usize) -> *mut c_void {
    let ptr = f!()(size);

    send_zone(ZoneKind::Malloc, ptr, size);

    ptr
});

shadow!(calloc(number: usize, size: usize) -> *mut c_void {
    let ptr = f!()(number, size);

    send_zone(ZoneKind::Calloc, ptr, number * size);

    ptr
});

shadow!(posix_memalign(ptr: *mut *mut c_void, alignment: usize, size: usize) -> c_int {
    let ret = f!()(ptr, alignment, size);

    if ret == 0 && let Some(ptr) = unsafe { ptr.as_mut() } {
        send_zone(ZoneKind::Memalign, unsafe { *ptr }, size);
    }

    ret
});

shadow!(memalign(alignment: usize, size: usize) -> *mut c_void {
    let ptr = f!()(alignment, size);

    send_zone(ZoneKind::Memalign, ptr, size);

    ptr
});

shadow!(aligned_alloc(alignment: usize, size: usize) -> *mut c_void {
    let ptr = f!()(alignment, size);

    send_zone(ZoneKind::Memalign, ptr, size);

    ptr
});

shadow!(pvalloc(size: usize) -> *mut c_void {
    let ptr = f!()(size);

    send_zone(ZoneKind::Memalign, ptr, size);

    ptr
});

shadow!(valloc(size: usize) -> *mut c_void {
    let ptr = f!()(size);

    send_zone(ZoneKind::Memalign, ptr, size);

    ptr
});

shadow!(realloc(old_ptr: *mut c_void, size: usize) -> *mut c_void {
    let new_ptr = f!()(old_ptr, size);

    remove_zone(old_ptr);
    if !new_ptr.is_null() {
        send_zone(ZoneKind::Realloc, new_ptr, size);
    }

    new_ptr
});

shadow!(reallocarray(old_ptr: *mut c_void, number: usize, size: usize) -> *mut c_void {
    let new_ptr = f!()(old_ptr, number, size);

    remove_zone(old_ptr);
    if !new_ptr.is_null() {
        send_zone(ZoneKind::Realloc, new_ptr, number * size);
    }

    new_ptr
});

shadow!(free(ptr: *mut c_void) -> () {
    f!()(ptr);

    remove_zone(ptr);
});

/*
 * Ensure we do not use our own shadowed functions within this library.
 * Recursion would be bad :)
 *
 * So we use a separate copy of malloc() and free() for our GlobalAlloc
 */

static RUST_MALLOC: AtomicFnPtr<extern "C" fn(usize) -> *mut c_void> =
    AtomicFnPtr::<extern "C" fn(usize) -> *mut c_void>::new(Some(c_stringify!("malloc")));
static RUST_FREE: AtomicFnPtr<extern "C" fn(*mut c_void)> =
    AtomicFnPtr::<extern "C" fn(*mut c_void)>::new(Some(c_stringify!("free")));

#[unsafe(no_mangle)]
pub fn rust_malloc(size: usize) -> *mut u8 {
    RUST_MALLOC.load_or_shadow()(size) as *mut u8
}

#[unsafe(no_mangle)]
pub fn rust_free(ptr: *mut u8) {
    RUST_FREE.load_or_shadow()(ptr as *mut c_void)
}

struct MyAllocator;
unsafe impl GlobalAlloc for MyAllocator {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        rust_malloc(layout.size())
    }

    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
        rust_free(ptr);
    }
}
#[global_allocator]
static GLOBAL_ALLOC: MyAllocator = MyAllocator;