peakmem-alloc 0.3.0

//!
//! An instrumenting allocator wrapper to compute (scoped) peak memory consumption.
//!
//! ## Example
//!
//! ```
//! use peakmem_alloc::*;
//! use std::alloc::System;
//!
//! #[global_allocator]
//! static GLOBAL: &PeakMemAlloc<System> = &INSTRUMENTED_SYSTEM;
//!
//! fn main() {
//!    GLOBAL.reset_peak_memory();
//!    let _x: Vec<u8> = Vec::with_capacity(1_024);
//!    println!(
//!        "Peak Memory used by function : {:#?}",
//!        GLOBAL.get_peak_memory()
//!    );
//! }
//! ```

#![deny(
    missing_debug_implementations,
    missing_copy_implementations,
    trivial_casts,
    trivial_numeric_casts,
    unused_import_braces,
    unused_imports,
    unused_qualifications,
    missing_docs
)]
#![cfg_attr(doc_cfg, feature(allocator_api))]
#![cfg_attr(doc_cfg, feature(doc_cfg))]

use std::{
    alloc::{GlobalAlloc, Layout, System},
    sync::atomic::{AtomicIsize, AtomicUsize, Ordering},
};

/// The PeakAllocTrait trait provides a common interface for all allocators.
///
/// This is mainly to allow for generic functions that can work with any allocator with type erasure.
///
pub trait PeakMemAllocTrait {
    /// Resets the peak memory to 0
    fn reset_peak_memory(&self);

    /// Get the peak memory consumption. This is the maximum that has been allocated since the last reset.
    ///
    /// Note that allocations of other threads may interfere with a measurement of a scope.
    fn get_peak_memory(&self) -> usize;
}

/// An allocator middleware which keeps track of peak memory consumption.
#[derive(Default, Debug)]
pub struct PeakMemAlloc<T: GlobalAlloc> {
    peak_bytes_allocated_tracker: AtomicIsize,
    peak_bytes_allocated: AtomicUsize,
    inner: T,
}

/// An instrumented instance of the system allocator.
pub static INSTRUMENTED_SYSTEM: PeakMemAlloc<System> = PeakMemAlloc {
    peak_bytes_allocated_tracker: AtomicIsize::new(0),
    peak_bytes_allocated: AtomicUsize::new(0),
    inner: System,
};

impl PeakMemAlloc<System> {
    /// Provides access to an instrumented instance of the system allocator.
    pub const fn system() -> Self {
        PeakMemAlloc {
            peak_bytes_allocated_tracker: AtomicIsize::new(0),
            peak_bytes_allocated: AtomicUsize::new(0),
            inner: System,
        }
    }
}

impl<T: GlobalAlloc> PeakMemAllocTrait for PeakMemAlloc<T> {
    /// Resets the peak memory to 0
    #[inline]
    fn reset_peak_memory(&self) {
        self.peak_bytes_allocated.store(0, Ordering::SeqCst);
        self.peak_bytes_allocated_tracker.store(0, Ordering::SeqCst);
    }

    /// Get the peak memory consumption
    #[inline]
    fn get_peak_memory(&self) -> usize {
        self.peak_bytes_allocated.load(Ordering::SeqCst)
    }
}
impl<T: GlobalAlloc> PeakMemAlloc<T> {
    /// Provides access to an instrumented instance of the given global
    /// allocator.
    pub const fn new(inner: T) -> Self {
        PeakMemAlloc {
            peak_bytes_allocated_tracker: AtomicIsize::new(0),
            peak_bytes_allocated: AtomicUsize::new(0),
            inner,
        }
    }

    #[inline]
    fn track_alloc(&self, bytes: usize) {
        let prev = self
            .peak_bytes_allocated_tracker
            .fetch_add(bytes as isize, Ordering::SeqCst);
        let current_peak = (prev + bytes as isize).max(0) as usize;
        self.peak_bytes_allocated
            .fetch_max(current_peak, Ordering::SeqCst);
    }

    #[inline]
    fn track_dealloc(&self, bytes: usize) {
        self.peak_bytes_allocated_tracker
            .fetch_sub(bytes as isize, Ordering::SeqCst);
    }
}

unsafe impl<'a, T: GlobalAlloc + 'a> GlobalAlloc for &'a PeakMemAlloc<T> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        (*self).alloc(layout)
    }

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

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

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

unsafe impl<T: GlobalAlloc> GlobalAlloc for PeakMemAlloc<T> {
    #[inline]
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        self.track_alloc(layout.size());
        self.inner.alloc(layout)
    }

    #[inline]
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        self.track_dealloc(layout.size());
        self.inner.dealloc(ptr, layout)
    }

    #[inline]
    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        self.track_alloc(layout.size());
        self.inner.alloc_zeroed(layout)
    }

    #[inline]
    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        match new_size.cmp(&layout.size()) {
            std::cmp::Ordering::Greater => {
                let difference = new_size - layout.size();
                self.track_alloc(difference);
            }
            std::cmp::Ordering::Less => {
                let difference = layout.size() - new_size;
                self.track_dealloc(difference);
            }
            _ => {}
        }

        self.inner.realloc(ptr, layout, new_size)
    }
}