heapster 0.4.0

//! **artisanal, pure-atomic heap telemetry for rust.**
//!
//! `heapster` is a stupid-lightweight, generic wrapper over any `GlobalAlloc`
//! that tracks allocations, deallocations, and reallocations using pure relaxed atomics.
//! it is designed to be always-on, allowing you to hunt down pathological memory patterns,
//! diff heap usage between code paths, and export raw allocator metrics to your telemetry
//! dashboards with minimal overhead.
//!
//! ## why heapster?
//!
//! heavyweight heap profilers are great for deep-dives, but they are often too slow
//! for production and too complex for simple CI assertions. `heapster` fills the gap:
//!
//! - **pure atomics**: no mutexes, no thread-locals, no external viewer files. just `AtomicUsize` with `Ordering::Relaxed`.
//! - **plug-and-play generic**: wrap `System`, `jemalloc`, `mimalloc`, or any custom allocator.
//! - **pathology hunting**: logarithmic size bucketing (histograms) tells you exactly what sizes are dominating your heap.
//! - **deep realloc tracking**: distinguishes between reallocations that grew in-place, shrank in-place, or forced a full memory copy.
//! - **zero-friction diffing**: exposes a [`Heapster::measure`] method to take clean snapshots of memory behavior around hot loops.
//!
//! ## quickstart
//!
//! wrap your global allocator of choice (e.g., `System`) in your `main.rs` or `lib.rs`:
//!
//! ```rust
//! use heapster::Heapster;
//! use std::alloc::System;
//!
//! #[global_allocator]
//! static GLOBAL: Heapster<System> = Heapster::new(System);
//!
//! fn main() {
//!     // ... do some heavy work ...
//!
//!     // see what has transpired in the heap
//!     let stats = GLOBAL.stats();
//!     println!("allocated: {} bytes", stats.alloc_sum);
//! }
//! ```
//!
//! ## measuring specific operations
//!
//! stop guessing if a change increased allocations. `heapster` lets you diff the heap stats
//! of critical sections of code using snapshot math:
//!
//! ```rust
//! # use heapster::Heapster;
//! # use std::alloc::System;
//! # #[global_allocator]
//! # static GLOBAL: Heapster<System> = Heapster::new(System);
//! let (result, heap_diff) = GLOBAL.measure(|| {
//!     // ... operation to measure ...
//!     42
//! });
//!
//! assert!(heap_diff.alloc_count < 10, "regression: the operation allocated too many times!");
//! ```

#![no_std]
#![deny(missing_docs)]
#![deny(clippy::all)]

#[cfg(feature = "fmt")]
mod fmt;
mod histogram;
mod stats;

pub use histogram::Histogram;
pub use stats::Stats;

#[cfg(feature = "fmt")]
extern crate alloc;

use core::{
    alloc::{GlobalAlloc, Layout},
    cmp,
    sync::atomic::{AtomicU64, AtomicUsize, Ordering},
};

static ALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);
static ALLOC_SUM: AtomicU64 = AtomicU64::new(0);
static ALLOC_BUCKETS: [AtomicUsize; 64] = [const { AtomicUsize::new(0) }; 64];
static ALLOC_FAIL_COUNT: AtomicUsize = AtomicUsize::new(0);

static DEALLOC_COUNT: AtomicUsize = AtomicUsize::new(0);
static DEALLOC_SUM: AtomicU64 = AtomicU64::new(0);

static REALLOC_GROWTH_COUNT: AtomicUsize = AtomicUsize::new(0);
static REALLOC_GROWTH_SUM: AtomicU64 = AtomicU64::new(0);
static REALLOC_GROWTH_BUCKETS: [AtomicUsize; 64] = [const { AtomicUsize::new(0) }; 64];

static REALLOC_SHRINK_COUNT: AtomicUsize = AtomicUsize::new(0);
static REALLOC_SHRINK_SUM: AtomicU64 = AtomicU64::new(0);
static REALLOC_SHRINK_BUCKETS: [AtomicUsize; 64] = [const { AtomicUsize::new(0) }; 64];

static REALLOC_MOVE_COUNT: AtomicUsize = AtomicUsize::new(0);
static REALLOC_MOVE_SUM: AtomicU64 = AtomicU64::new(0);
static REALLOC_FAIL_COUNT: AtomicUsize = AtomicUsize::new(0);

static USE_CURR: AtomicUsize = AtomicUsize::new(0);
static USE_MAX: AtomicUsize = AtomicUsize::new(0);

/// A global allocator enhanced with stats.
#[derive(Debug, Default, Clone, Copy)]
pub struct Heapster<A: GlobalAlloc>(A);

fn bucket_snapshot(buckets: &[AtomicUsize; 64]) -> Histogram {
    let mut out = [0usize; 64];
    for (i, b) in buckets.iter().enumerate() {
        out[i] = b.load(Ordering::Relaxed);
    }
    Histogram { buckets: out }
}

impl<A: GlobalAlloc> Heapster<A> {
    /// Wraps an allocator, facilitating useful stats.
    pub const fn new(alloc: A) -> Self {
        Self(alloc)
    }

    /// Returns a reference to the underlying allocator.
    pub const fn inner(&self) -> &A {
        &self.0
    }

    /// Returns the total number of allocations.
    pub fn alloc_count(&self) -> usize {
        ALLOC_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the sum of all allocations.
    pub fn alloc_sum(&self) -> u64 {
        ALLOC_SUM.load(Ordering::Relaxed)
    }

    /// Returns a histogram representing the number
    /// of allocations of different sizes.
    pub fn alloc_histogram(&self) -> Histogram {
        bucket_snapshot(&ALLOC_BUCKETS)
    }

    /// Returns the total number of failed allocations.
    pub fn alloc_fail_count(&self) -> usize {
        ALLOC_FAIL_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the total number of deallocations.
    pub fn dealloc_count(&self) -> usize {
        DEALLOC_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the sum of all deallocations.
    pub fn dealloc_sum(&self) -> u64 {
        DEALLOC_SUM.load(Ordering::Relaxed)
    }

    /// Returns the total number of reallocations caused by object growth.
    pub fn realloc_growth_count(&self) -> usize {
        REALLOC_GROWTH_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the sum of all reallocations caused by object growth.
    pub fn realloc_growth_sum(&self) -> u64 {
        REALLOC_GROWTH_SUM.load(Ordering::Relaxed)
    }

    /// Returns a histogram representing the number
    /// of growth reallocations of different sizes.
    pub fn realloc_growth_histogram(&self) -> Histogram {
        bucket_snapshot(&REALLOC_GROWTH_BUCKETS)
    }

    /// Returns the total number of reallocations caused by object shrinkage.
    pub fn realloc_shrink_count(&self) -> usize {
        REALLOC_SHRINK_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the sum of all reallocations caused by object shrinkage.
    pub fn realloc_shrink_sum(&self) -> u64 {
        REALLOC_SHRINK_SUM.load(Ordering::Relaxed)
    }

    /// Returns a histogram representing the number
    /// of shrink reallocations of different sizes.
    pub fn realloc_shrink_histogram(&self) -> Histogram {
        bucket_snapshot(&REALLOC_SHRINK_BUCKETS)
    }

    /// Returns the total number of full reallocations.
    pub fn realloc_move_count(&self) -> usize {
        REALLOC_MOVE_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the sum of all full reallocations.
    pub fn realloc_move_sum(&self) -> u64 {
        REALLOC_MOVE_SUM.load(Ordering::Relaxed)
    }

    /// Returns the total number of failed reallocations.
    pub fn realloc_fail_count(&self) -> usize {
        REALLOC_FAIL_COUNT.load(Ordering::Relaxed)
    }

    /// Returns the average size of allocations.
    pub fn alloc_avg(&self) -> Option<usize> {
        let sum = ALLOC_SUM.load(Ordering::Relaxed);
        let count = ALLOC_COUNT.load(Ordering::Relaxed);
        sum.checked_div(count as u64).map(|avg| avg as usize)
    }

    /// Returns the average size of deallocations.
    pub fn dealloc_avg(&self) -> Option<usize> {
        let sum = DEALLOC_SUM.load(Ordering::Relaxed);
        let count = DEALLOC_COUNT.load(Ordering::Relaxed);
        sum.checked_div(count as u64).map(|avg| avg as usize)
    }

    /// Returns the average size of reallocations caused by object growth.
    pub fn realloc_growth_avg(&self) -> Option<usize> {
        let sum = REALLOC_GROWTH_SUM.load(Ordering::Relaxed);
        let count = REALLOC_GROWTH_COUNT.load(Ordering::Relaxed);
        sum.checked_div(count as u64).map(|avg| avg as usize)
    }

    /// Returns the average size of reallocations caused by object shrinkage.
    pub fn realloc_shrink_avg(&self) -> Option<usize> {
        let sum = REALLOC_SHRINK_SUM.load(Ordering::Relaxed);
        let count = REALLOC_SHRINK_COUNT.load(Ordering::Relaxed);
        sum.checked_div(count as u64).map(|avg| avg as usize)
    }

    /// Returns the average size of full reallocations.
    pub fn realloc_move_avg(&self) -> Option<usize> {
        let sum = REALLOC_MOVE_SUM.load(Ordering::Relaxed);
        let count = REALLOC_MOVE_COUNT.load(Ordering::Relaxed);
        sum.checked_div(count as u64).map(|avg| avg as usize)
    }

    /// Returns current heap use.
    pub fn use_curr(&self) -> usize {
        USE_CURR.load(Ordering::Relaxed)
    }

    /// Returns maximum recorded heap use.
    pub fn use_max(&self) -> usize {
        USE_MAX.load(Ordering::Relaxed)
    }

    /// Measures the heap stats for the given operation, returning its
    /// result alongside the [`Stats`] object.
    pub fn measure<R>(&self, f: impl FnOnce() -> R) -> (R, Stats) {
        let before = self.stats();
        let r = f();
        let after = self.stats();
        (r, &after - &before)
    }

    /// Sets the stats to 0, except for current heap use (which is unaffected)
    /// and maximum heap use, which is reset to the value of current heap use.
    ///
    /// **Concurrency note:** `reset` is not synchronized with allocator activity.
    /// If another thread is mid-allocation when `reset` runs, its increment may
    /// land on a freshly-zeroed counter, briefly producing skewed values (or, in
    /// rare cases, a transient apparent decrease in `use_curr`). For measuring a
    /// specific operation in a multi-threaded program, prefer [`Heapster::measure`],
    /// which uses snapshot diffing and avoids touching shared state.
    pub fn reset(&self) {
        ALLOC_SUM.store(0, Ordering::Relaxed);
        ALLOC_COUNT.store(0, Ordering::Relaxed);
        for b in &ALLOC_BUCKETS {
            b.store(0, Ordering::Relaxed);
        }
        ALLOC_FAIL_COUNT.store(0, Ordering::Relaxed);

        DEALLOC_SUM.store(0, Ordering::Relaxed);
        DEALLOC_COUNT.store(0, Ordering::Relaxed);

        REALLOC_GROWTH_COUNT.store(0, Ordering::Relaxed);
        REALLOC_GROWTH_SUM.store(0, Ordering::Relaxed);
        for b in &REALLOC_GROWTH_BUCKETS {
            b.store(0, Ordering::Relaxed);
        }

        REALLOC_SHRINK_COUNT.store(0, Ordering::Relaxed);
        REALLOC_SHRINK_SUM.store(0, Ordering::Relaxed);
        for b in &REALLOC_SHRINK_BUCKETS {
            b.store(0, Ordering::Relaxed);
        }

        REALLOC_MOVE_COUNT.store(0, Ordering::Relaxed);
        REALLOC_MOVE_SUM.store(0, Ordering::Relaxed);
        REALLOC_FAIL_COUNT.store(0, Ordering::Relaxed);

        USE_MAX.store(self.use_curr(), Ordering::Relaxed);
    }
}

#[inline]
fn bucket_of(size: usize) -> usize {
    debug_assert!(size > 0);
    (usize::BITS - 1 - size.leading_zeros()) as usize
}

unsafe impl<A: GlobalAlloc> GlobalAlloc for Heapster<A> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let ret = unsafe { self.0.alloc(layout) };
        if !ret.is_null() {
            let size = layout.size();
            ALLOC_SUM.fetch_add(size as u64, Ordering::Relaxed);
            ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
            let curr = USE_CURR.fetch_add(size, Ordering::Relaxed) + size;
            USE_MAX.fetch_max(curr, Ordering::Relaxed);
            ALLOC_BUCKETS[bucket_of(size)].fetch_add(1, Ordering::Relaxed);
        } else {
            ALLOC_FAIL_COUNT.fetch_add(1, Ordering::Relaxed);
        }

        ret
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        unsafe { self.0.dealloc(ptr, layout) };
        let size = layout.size();
        USE_CURR.fetch_sub(size, Ordering::Relaxed);
        DEALLOC_SUM.fetch_add(size as u64, Ordering::Relaxed);
        DEALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        let new_ptr = unsafe { self.0.realloc(ptr, layout, new_size) };
        if !new_ptr.is_null() {
            if new_size >= layout.size() {
                let diff = new_size - layout.size();
                REALLOC_GROWTH_COUNT.fetch_add(1, Ordering::Relaxed);
                REALLOC_GROWTH_SUM.fetch_add(diff as u64, Ordering::Relaxed);
                let curr = USE_CURR.fetch_add(diff, Ordering::Relaxed) + diff;
                USE_MAX.fetch_max(curr, Ordering::Relaxed);
                REALLOC_GROWTH_BUCKETS[bucket_of(diff)].fetch_add(1, Ordering::Relaxed);
            } else {
                let diff = layout.size() - new_size;
                REALLOC_SHRINK_COUNT.fetch_add(1, Ordering::Relaxed);
                REALLOC_SHRINK_SUM.fetch_add(diff as u64, Ordering::Relaxed);
                USE_CURR.fetch_sub(diff, Ordering::Relaxed);
                REALLOC_SHRINK_BUCKETS[bucket_of(diff)].fetch_add(1, Ordering::Relaxed);
            }
            if new_ptr != ptr {
                REALLOC_MOVE_COUNT.fetch_add(1, Ordering::Relaxed);
                REALLOC_MOVE_SUM
                    .fetch_add(cmp::min(layout.size(), new_size) as u64, Ordering::Relaxed);
            }
        } else {
            REALLOC_FAIL_COUNT.fetch_add(1, Ordering::Relaxed);
        }

        new_ptr
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        let ret = unsafe { self.0.alloc_zeroed(layout) };
        if !ret.is_null() {
            let size = layout.size();
            ALLOC_SUM.fetch_add(size as u64, Ordering::Relaxed);
            ALLOC_COUNT.fetch_add(1, Ordering::Relaxed);
            let curr = USE_CURR.fetch_add(size, Ordering::Relaxed) + size;
            USE_MAX.fetch_max(curr, Ordering::Relaxed);
            ALLOC_BUCKETS[bucket_of(size)].fetch_add(1, Ordering::Relaxed);
        } else {
            ALLOC_FAIL_COUNT.fetch_add(1, Ordering::Relaxed);
        }

        ret
    }
}