rs_malloc_tracker 1.0.1

#![allow(clippy::deref_addrof)] // necessary for static muts

use std::{
    ffi::c_void,
    sync::{
        LazyLock, Mutex, RwLock,
        mpsc::{Sender, channel},
    },
    thread::{self, JoinHandle},
};

use ctor::{ctor, dtor};
#[cfg(feature = "usable_size")]
use libc::malloc_usable_size;
use log::*;
use rustc_hash::FxHashMap;

use crate::{RUST_FREE, RUST_MALLOC, get_frame, thread_local_sender::ThreadLocalSender};

#[derive(Copy, Clone, Debug, PartialEq)]
pub enum ZoneKind {
    Malloc,
    Calloc,
    Memalign,
    Realloc,
}

#[allow(dead_code)]
#[derive(Clone)]
pub struct MemoryZone {
    pub size: usize,
    pub usable_size: usize,
    pub ip: Option<usize>,
    pub kind: ZoneKind,
}

#[unsafe(no_mangle)]
pub static MEMORY_ZONES: LazyLock<RwLock<FxHashMap<usize, MemoryZone>>> =
    LazyLock::new(|| RwLock::new(FxHashMap::default()));
#[unsafe(no_mangle)]
pub static mut SEND_FAILS: LazyLock<RwLock<FxHashMap<String, usize>>> =
    LazyLock::new(|| RwLock::new(FxHashMap::default()));

// send_zone -> *alloc -> [frame we want]
#[cfg(debug_assertions)]
const FRAME: i32 = 3;
#[cfg(not(debug_assertions))]
const FRAME: i32 = 3;

#[inline(never)]
pub fn send_zone(kind: ZoneKind, ptr: *mut c_void, size: usize) {
    if RUST_MALLOC.try_load().is_none() || RUST_FREE.try_load().is_none() {
        return;
    }

    let ip = get_frame!(FRAME).map(|v| v as usize);

    #[cfg(feature = "usable_size")]
    let usable_size = unsafe { malloc_usable_size(ptr) };
    #[cfg(not(feature = "usable_size"))]
    let usable_size = 0usize;

    ThreadLocalSender::new(&SENDER)
        .try_send(&move || {
            Operation::Alloc((
                ptr as usize,
                MemoryZone {
                    size,
                    usable_size,
                    ip,
                    kind,
                },
            ))
        })
        .inspect_err(|e| unsafe {
            *(*&raw mut SEND_FAILS)
                .write()
                .unwrap()
                .entry(e.clone())
                .or_default() += 1;
        })
        .unwrap_or(());
}

pub fn remove_zone(ptr: *mut c_void) {
    ThreadLocalSender::new(&SENDER)
        .try_send(&move || Operation::Free(ptr as *mut _ as usize))
        .inspect_err(|e| unsafe {
            *(*&raw mut SEND_FAILS)
                .write()
                .unwrap()
                .entry(e.clone())
                .or_default() += 1;
        })
        .unwrap_or(());
}

#[cfg(test)]
pub fn test_wait_processed() {
    let (sender, receiver) = channel::<()>();
    SENDER
        .lock()
        .unwrap()
        .as_ref()
        .unwrap()
        .send(Operation::TestFence(sender))
        .unwrap();
    receiver.recv().unwrap();
}

pub enum Operation {
    Alloc((usize, MemoryZone)),
    Free(usize),
    #[cfg(test)]
    TestFence(Sender<()>), // used by tests to wait for pending operations to be processed
    Shutdown,
}

static SENDER: LazyLock<Mutex<Option<Sender<Operation>>>> = LazyLock::new(|| Mutex::new(None));
static RECEIVER_THREAD: LazyLock<Mutex<Option<JoinHandle<()>>>> =
    LazyLock::new(|| Mutex::new(None));

#[ctor]
fn init() {
    let (sender, receiver) = channel::<Operation>();

    SENDER.lock().unwrap().replace(sender);

    RECEIVER_THREAD
        .lock()
        .unwrap()
        .replace(thread::spawn(move || {
            debug!("receiver thread starting");

            loop {
                match receiver.recv() {
                    Ok(Operation::Alloc((addr, zone))) => {
                        MEMORY_ZONES.write().unwrap().insert(addr, zone);
                    }
                    Ok(Operation::Free(zone)) => {
                        MEMORY_ZONES.write().unwrap().remove(&zone);
                    }
                    #[cfg(test)]
                    Ok(Operation::TestFence(s)) => s.send(()).unwrap(),
                    Ok(Operation::Shutdown) => break,
                    Err(_) => break,
                }
            }

            debug!("receiver thread disconnecting");
            SENDER.lock().unwrap().take();
        }));
}

#[dtor]
fn deinit() {
    SENDER
        .lock()
        .unwrap()
        .as_ref()
        .map(|s| s.send(Operation::Shutdown));
    RECEIVER_THREAD.lock().unwrap().take();
}