lib.rs

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
use std::cmp::{Ordering, PartialOrd, Ord};
use std::collections::{HashMap, hash_map::Entry};
use std::sync::{LazyLock, Mutex};
//use sorted_vec::SortedSet;

// cargo test -- --nocapture --test-threads=1

const MAX_TRACKED_ALLOCATIONS: usize = 2048;//1024;
const MAX_FRAMES: usize = 50;
const SKIPPED_FRAMES: usize = 5;
const MIN_SIZE: usize = 512_000;//4096;

thread_local! {
    static IN_ALLOCATION: Cell<bool> = Cell::new(false);
}

#[derive(PartialEq, Eq)]
struct AllocSite {
    frames: [*mut std::ffi::c_void; MAX_FRAMES],
    ptr: *mut u8,
    size: usize,
    noted: bool,
}

unsafe impl Send for AllocSite {}

impl Ord for AllocSite {
    fn cmp(&self, other: &Self) -> Ordering {
	self.partial_cmp(other).unwrap()
    }
}

impl PartialOrd for AllocSite {
    fn partial_cmp(&self, other: &AllocSite) -> Option<Ordering> {
	Some(match self.size.cmp(&other.size) {
	    Ordering::Equal => self.frames.cmp(&other.frames),
	    o => o,
	})
    }
}

static LARGEST_ALLOCATION_SITES: LazyLock<Mutex<HashMap<usize, AllocSite>>> = LazyLock::new(Default::default);

#[derive(Default)]
pub struct AccountingAlloc<A = System> {
    allocator: A,
}

impl AccountingAlloc<System> {
    pub const fn new() -> Self {
        Self::with_allocator(System)
    }
}

impl<A> AccountingAlloc<A> {
    pub const fn with_allocator(allocator: A) -> Self {
        Self {
	    allocator,
	}
    }

    fn remove_allocation(&self, ptr: *mut u8, size: usize) {
        if size < MIN_SIZE {
            return;
        }
	let old = IN_ALLOCATION.with(|status| {
	    status.replace(true)
	});
	if old {
	    return;
	}
        let mut sites = LARGEST_ALLOCATION_SITES.lock().unwrap();
        match sites.entry(ptr as usize) {
            Entry::Occupied(e) => {
                e.remove();
            }
            _ => {}
        }
        IN_ALLOCATION.with(|status| status.set(old));
    }

    fn record_allocation(&self, ptr: *mut u8, size: usize) {
	if size < MIN_SIZE {
	    return;
	}
	let old = IN_ALLOCATION.with(|status| {
	    status.replace(true)
	});
	if old {
	    return;
	}
	let mut num_skipped = 0;
	let mut num_frames = 0;
	let mut frames = [std::ptr::null_mut(); MAX_FRAMES];
	backtrace::trace(|frame| {
	    if num_skipped < SKIPPED_FRAMES {
		num_skipped += 1;
		return true;
	    }
	    if num_frames >= MAX_FRAMES {
		return false;
	    }
	    frames[num_frames] = frame.ip();
	    num_frames += 1;
	    true
	});
	let site = AllocSite {
	    frames,
	    size,
            ptr,
            noted: false,
	};
	let mut sites = LARGEST_ALLOCATION_SITES.lock().unwrap();

	if sites.len() < MAX_TRACKED_ALLOCATIONS {
	    sites.insert(ptr as usize, site);
        } else if let Some(key) = sites.iter().find(|(_, s)| s.size < site.size).map(|(k, _)| k.clone()) {
            sites.remove(&key);
            sites.insert(ptr as usize, site);
        }
	/*} else if site.size > sites.first().map(|site| site.size).unwrap_or_default() {
	    sites.remove_index(0);
	    sites.find_or_push(site);
	}*/

	IN_ALLOCATION.with(|status| status.set(old));
    }

    pub fn note_allocation(&self, ptr: *mut u8, size: usize) {
        if size < MIN_SIZE {
            return;
        }
	IN_ALLOCATION.with(|status| status.set(true));
        if let Some(site) = LARGEST_ALLOCATION_SITES.lock().unwrap().get_mut(&(ptr as usize)) {
            site.noted = true;
        }
	IN_ALLOCATION.with(|status| status.set(false));
    }

    pub fn dump_unmeasured_allocations(&self) {
	IN_ALLOCATION.with(|status| status.set(true));
	let sites = LARGEST_ALLOCATION_SITES.lock().unwrap();
	let default = "???".to_owned();
	for site in sites.values().filter(|site| !site.noted) {
	    let mut resolved = vec![];
	    for ip in site.frames.iter().filter(|ip| !ip.is_null()) {
		backtrace::resolve(*ip, |symbol| {
		    resolved.push((
			symbol.filename().map(|f| f.to_owned()),
			symbol.lineno(),
			symbol.name().map(|n| format!("{}", n)),
		    ));
		});
	    };

	    println!("---\n{}\n", site.size);
	    for (filename, line, symbol) in &resolved {
		let fname = filename.as_ref().map(|f| f.display().to_string());
		println!(
		    "{}:{} ({})",
		    fname.as_ref().unwrap_or(&default),
		    line.unwrap_or_default(),
		    symbol.as_ref().unwrap_or(&default),
		);
	    }
	}
	IN_ALLOCATION.with(|status| status.set(false));
    }

    pub fn dump_largest_allocations(&self) {
	IN_ALLOCATION.with(|status| status.set(true));
	let sites = LARGEST_ALLOCATION_SITES.lock().unwrap();
	let default = "???".to_owned();
	for site in sites.values() {
	    let mut resolved = vec![];
	    for ip in site.frames.iter().filter(|ip| !ip.is_null()) {
		backtrace::resolve(*ip, |symbol| {
		    resolved.push((
			symbol.filename().map(|f| f.to_owned()),
			symbol.lineno(),
			symbol.name().map(|n| format!("{}", n)),
		    ));
		});
	    };

	    println!("---\n{}\n", site.size);
	    for (filename, line, symbol) in &resolved {
		let fname = filename.as_ref().map(|f| f.display().to_string());
		println!(
		    "{}:{} ({})",
		    fname.as_ref().unwrap_or(&default),
		    line.unwrap_or_default(),
		    symbol.as_ref().unwrap_or(&default),
		);
	    }
	}
	IN_ALLOCATION.with(|status| status.set(false));
    }
}

unsafe impl<A: GlobalAlloc> GlobalAlloc for AccountingAlloc<A> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ptr = self.allocator.alloc(layout);
        self.record_allocation(ptr, layout.size());
        ptr
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.allocator.dealloc(ptr, layout);
        self.remove_allocation(ptr, layout.size());
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ptr = self.allocator.alloc_zeroed(layout);
        self.record_allocation(ptr, layout.size());
        ptr
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        self.remove_allocation(ptr, layout.size());
        let ptr = self.allocator.realloc(ptr, layout, new_size);
        self.record_allocation(ptr, new_size);
        ptr
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::convert::identity;
    use std::thread::scope;

    #[derive(Default)]
    /// A fake [`GlobalAlloc`] memory allocator for testing, which doesn't actually allocate the memory requested.
    ///
    /// This allocator must not actually be registered as the global allocator using `#[global_allocator]`. The pointers
    /// returned by this allocator should be treated as opaque and only used in subsequent calls to `<TestAlloc as
    /// GlobalAlloc>` methods.
    ///
    /// Though technically panicking from [`GlobalAlloc`] methods is currently considered UB, this is only relevant if
    /// the [`GlobalAlloc`] is actually registered as the global allocator. Thus, to verify correctness during tests,
    /// `dealloc` and `realloc` will panic if `layout` is not equal to the `layout` provided to `alloc`.
    struct TestAlloc;

    /// Metadata for an allocation made by [`TestAlloc`].
    struct Allocation {
        layout: Layout,
    }

    /// A handle to an allocation made by an `AccountingAlloc<TestAlloc>`.
    ///
    /// The allocation will be deallocated when this structure is dropped.
    struct AllocationHandle<'a> {
        allocator: &'a AccountingAlloc<TestAlloc>,
        ptr: *mut u8,
        layout: Layout,
    }

    /// Make some variously sized test allocations in separate threads using the provided `allocate` function and call
    /// `callback`.
    ///
    /// The `AllocationHandle`s returned by `allocate` are passed to `callback`. After `callback` returns, all spawned
    /// threads are waited on to terminate, and the return value of `callback` is returned.
    fn test_allocations<'a, T>(
        allocator: &'a AccountingAlloc<TestAlloc>,
        allocate: fn(&'a AccountingAlloc<TestAlloc>, Layout) -> AllocationHandle<'a>,
        callback: impl FnOnce(Vec<AllocationHandle<'a>>) -> T,
    ) -> T {
        let layouts: Vec<_> = (1..10).map(|idx| Layout::array::<u8>(10000 * idx).unwrap()).collect();
        let (allocations_tx, allocations_rx) = std::sync::mpsc::channel();
        scope(|scope| {
            for layout in layouts.clone() {
                let allocations_tx = allocations_tx.clone();
                scope.spawn(move || allocations_tx.send(allocate(allocator, layout)).unwrap());
            }
            drop(allocations_tx);
            callback(allocations_rx.into_iter().collect())
        })
    }

    #[test]
    fn alloc() {
        let allocator = Default::default();
        let /*(_allocations, expected)*/_ = test_allocations(&allocator, AllocationHandle::new, |_allocations| {
            // test `allocator.count()` while threads are still alive.
            /*let expected = expected_counts(&allocations);
            assert_eq!(allocator.count(), expected);
            (allocations, expected)*/
	    ()
        });

	allocator.dump_largest_allocations();
        // test `allocator.count()` again after the threads are dead.
        /*assert_eq!(
            allocator.count(),
            AllocStats { since_last: Default::default(), ..expected }
        );*/
    }

    #[test]
    fn alloc_zeroed() {
        let allocator = &Default::default();
        let _ = test_allocations(&allocator, AllocationHandle::new_zeroed, identity);
	allocator.dump_largest_allocations();
    }

    #[test]
    fn realloc() {
        let allocator = &Default::default();
        let mut allocations = test_allocations(&allocator, AllocationHandle::new, identity);

        scope(|scope| {
            for allocation in &mut allocations {
                scope.spawn(move || allocation.realloc(allocation.layout.size() * 2));
            }
        });

	allocator.dump_largest_allocations();
    }

    unsafe impl GlobalAlloc for TestAlloc {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            // Since we're not actually accessing this allocated memory in tests, we don't actually need to allocate
            // anything, but allocating something using the system allocator lets us leverage it for double-free and
            // leak detection. We save the Layout in a Box, however, to be able to later assert that the Layout upon
            // dealloc/realloc matches, for correctness verification.
            Box::into_raw(Box::new(Allocation { layout })) as *mut u8
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            // Claim ownership of the allocation and free it.
            let allocation = Box::from_raw(ptr as *mut Allocation);
            assert_eq!(layout, allocation.layout);
            drop(allocation);
        }

        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            self.alloc(layout)
        }

        unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
            self.dealloc(ptr, layout);
            self.alloc(Layout::from_size_align_unchecked(new_size, layout.align()))
        }
    }

    impl<'a> AllocationHandle<'a> {
        /// Make a new allocation on `allocator` with the given `layout`.
        fn new(allocator: &'a AccountingAlloc<TestAlloc>, layout: Layout) -> Self {
            Self { allocator, ptr: unsafe { allocator.alloc(layout) }, layout }
        }

        /// Make a new zeroed allocation on `allocator` with the given `layout`.
        fn new_zeroed(allocator: &'a AccountingAlloc<TestAlloc>, layout: Layout) -> Self {
            Self { allocator, ptr: unsafe { allocator.alloc_zeroed(layout) }, layout }
        }

        /// Resize an allocation.
        fn realloc(&mut self, new_size: usize) {
            unsafe {
                self.ptr = self.allocator.realloc(self.ptr, self.layout, new_size);
                self.layout = Layout::from_size_align_unchecked(new_size, self.layout.align());
            }
        }
    }

    unsafe impl Send for AllocationHandle<'_> {}

    impl Drop for AllocationHandle<'_> {
        fn drop(&mut self) {
            unsafe { self.allocator.dealloc(self.ptr, self.layout) };
        }
    }
}