rs_malloc_tracker 1.0.1

/// Implements an atomic function pointer type (used by `shadow!`)
#[macro_export]
macro_rules! fn_impl {
    ($ftype: ty) => {
        /*
         * The reason we are doing this weird thing with a static item and function
         * declarations is that function pointers *can't* be generics.
         * And I don't want to use a pointer-to-a-function-pointer as that would require an
         * allocation to store it.
         */
        impl $crate::AtomicFnPtr<$ftype> {
            pub const fn new(shadow_name: Option<&'static std::ffi::CStr>) -> Self {
                Self {
                    ptr: $crate::AtomicUsize::new(0),
                    shadow_name,
                    _t: std::marker::PhantomData,
                }
            }

            /*
             * The reasons we can avoid using a Mutex are:
             *
             * 1. We're just storing a function pointer which is representable as usize with some
             *    transmutation
             * 2. If we concurrently initialize this value twice, no big deal; dlsym() is cheap.
             */
            pub fn try_load(&self) -> Option<$ftype> {
                let ptr = self.ptr.load(std::sync::atomic::Ordering::Acquire) as *mut c_void;

                if ptr.is_null() {
                    None
                } else {
                    Some(unsafe { std::mem::transmute::<*mut c_void, $ftype>(ptr) })
                }
            }

            #[allow(dead_code)]
            fn load_or_shadow(&self) -> $ftype {
                if let Some(f) = self.try_load() {
                    f
                } else {
                    let ptr = unsafe {
                        libc::dlsym(
                            libc::RTLD_NEXT,
                            self.shadow_name.expect("Missing name").as_ptr(),
                        )
                    };
                    let ptr: $ftype = unsafe { std::mem::transmute(ptr) };
                    self.store(Some(ptr));
                    self.try_load().unwrap()
                }
            }

            pub fn store(&self, val: Option<$ftype>) {
                if let Some(val) = val {
                    let val = val as *const ();
                    self.ptr.store(
                        val as *const _ as usize,
                        std::sync::atomic::Ordering::Release,
                    )
                } else {
                    self.ptr.store(0usize, std::sync::atomic::Ordering::Release);
                }
            }
        }
    };
}

/// Shadows an existing symbol. A pointer to the shadowed function is available by calling `f!()`.
/// The function signature must be declared using `fn_impl!`.
///
/// Example:
///
/// ```rust
/// fn_impl!(extern "C" fn (usize) -> *mut c_void);
///
/// shadow!(malloc(size: usize) -> *mut c_void {
///     let ptr = f!()(size);
///
///     send_zone(ZoneKind::Malloc, ptr, size);
///
///     ptr
/// });
/// ```
#[macro_export]
macro_rules! shadow {
    ($name: ident ($($v:ident: $t:ty),+) -> $return: ty $block: block) => {
        macro_rules! ftype {
            () => { extern "C" fn ($($t),+) -> $return }
        }

        /// # Safety
        ///
        /// Same rules as for malloc(2) apply.
        #[unsafe(no_mangle)]
        #[allow(unused)]
        pub unsafe extern "C" fn $name($($v: $t),+) -> $return {
            static SHADOWED_ITEM: $crate::AtomicFnPtr<ftype!()> = $crate::AtomicFnPtr::<ftype!()>::new(Some($crate::c_stringify!(stringify!($name))));

            let f: ftype!() = if let Some(f) = SHADOWED_ITEM.try_load() {
                f
            } else {
                let ptr = unsafe { libc::dlsym(libc::RTLD_NEXT, SHADOWED_ITEM.shadow_name.expect("Missing name").as_ptr()) };
                let ptr = unsafe { std::mem::transmute::<*mut c_void, ftype!()>(ptr) };
                SHADOWED_ITEM.store(Some(ptr));
                SHADOWED_ITEM.try_load().unwrap()
            };

            macro_rules! f {
                () => {
                    (f as ftype!())
                }
            }

            $block
        }
    };
}

/// Get a stack pointer to a single frame.
///
/// We could use [`backtrace-rs`](https://github.com/rust-lang/backtrace-rs). However I noticed two
/// problems with that:
/// 1. It is slower (iterates over all frames regardless of whether it is going to be useful). This
///    is a problem in our use-case since we run on the hot path for malloc() !
/// 2. I ran into a weird deadlock where `_Unwind_Backtrace` was stuck on
///    `_dl_find_object_to_external`? related to 1. I don't think this should even /attempt/ to
///    resolve.
///
/// This has to be a macro as to not do heap allocations.
///
/// Per the [libunwind
/// documentation](https://www.nongnu.org/libunwind/man/libunwind(3).html#section_5):
/// > All libunwind routines are thread-safe. What this means is that multiple threads may use libunwind simulatenously. However, any given cursor may be accessed by only one thread at any given time.
#[macro_export]
macro_rules! get_frame {
    ($n: expr) => {{
        unsafe extern "C" {
            fn unw_backtrace(buffer: *mut *mut c_void, size: c_int) -> c_int;
        }

        use std::ffi::{c_int, c_void};

        let mut buf = [std::ptr::null_mut() as *mut c_void; $n as usize];
        let resolved = unsafe { unw_backtrace(buf.as_mut_ptr(), $n as i32) };

        if resolved == $n {
            buf.last().map(|v| *v as *mut c_void)
        } else {
            None
        }
    }};
}