#![cfg_attr(docsrs, feature(doc_cfg))]
#![deny(missing_docs)]
#![deny(clippy::pedantic)]
#![allow(clippy::inline_always)]
#![allow(clippy::module_name_repetitions)]
use std::{
error, fmt,
sync::{
atomic::{AtomicBool, AtomicUsize, Ordering},
Arc,
},
};
mod allocator;
mod stack;
mod token;
#[cfg(feature = "tracing-compat")]
mod tracing;
mod util;
use token::with_suspended_allocation_group;
pub use crate::allocator::Allocator;
pub use crate::token::{AllocationGroupId, AllocationGroupToken, AllocationGuard};
#[cfg(feature = "tracing-compat")]
pub use crate::tracing::AllocationLayer;
static TRACKING_ENABLED: AtomicBool = AtomicBool::new(false);
static mut GLOBAL_TRACKER: Option<Tracker> = None;
static GLOBAL_INIT: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;
pub trait AllocationTracker {
fn allocated(
&self,
addr: usize,
object_size: usize,
wrapped_size: usize,
group_id: AllocationGroupId,
);
fn deallocated(
&self,
addr: usize,
object_size: usize,
wrapped_size: usize,
source_group_id: AllocationGroupId,
current_group_id: AllocationGroupId,
);
}
struct Tracker {
tracker: Arc<dyn AllocationTracker + Send + Sync + 'static>,
}
impl Tracker {
fn from_allocation_tracker<T>(allocation_tracker: T) -> Self
where
T: AllocationTracker + Send + Sync + 'static,
{
Self {
tracker: Arc::new(allocation_tracker),
}
}
fn allocated(
&self,
addr: usize,
object_size: usize,
wrapped_size: usize,
group_id: AllocationGroupId,
) {
self.tracker
.allocated(addr, object_size, wrapped_size, group_id);
}
fn deallocated(
&self,
addr: usize,
object_size: usize,
wrapped_size: usize,
source_group_id: AllocationGroupId,
current_group_id: AllocationGroupId,
) {
self.tracker.deallocated(
addr,
object_size,
wrapped_size,
source_group_id,
current_group_id,
);
}
}
#[derive(Debug)]
pub struct SetTrackerError {
_sealed: (),
}
impl fmt::Display for SetTrackerError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("a global tracker has already been set")
}
}
impl error::Error for SetTrackerError {}
pub struct AllocationRegistry;
impl AllocationRegistry {
pub fn enable_tracking() {
TRACKING_ENABLED.store(true, Ordering::SeqCst);
}
pub fn disable_tracking() {
TRACKING_ENABLED.store(false, Ordering::SeqCst);
}
pub fn set_global_tracker<T>(tracker: T) -> Result<(), SetTrackerError>
where
T: AllocationTracker + Send + Sync + 'static,
{
if GLOBAL_INIT
.compare_exchange(
UNINITIALIZED,
INITIALIZING,
Ordering::AcqRel,
Ordering::Relaxed,
)
.is_ok()
{
unsafe {
GLOBAL_TRACKER = Some(Tracker::from_allocation_tracker(tracker));
}
GLOBAL_INIT.store(INITIALIZED, Ordering::Release);
Ok(())
} else {
Err(SetTrackerError { _sealed: () })
}
}
pub fn untracked<F, R>(f: F) -> R
where
F: FnOnce() -> R,
{
with_suspended_allocation_group(f)
}
#[doc(hidden)]
pub unsafe fn clear_global_tracker() {
GLOBAL_INIT.store(INITIALIZING, Ordering::Release);
GLOBAL_TRACKER = None;
GLOBAL_INIT.store(UNINITIALIZED, Ordering::Release);
}
}
#[inline(always)]
fn get_global_tracker() -> Option<&'static Tracker> {
if !TRACKING_ENABLED.load(Ordering::Relaxed) {
return None;
}
if GLOBAL_INIT.load(Ordering::Acquire) != INITIALIZED {
return None;
}
unsafe {
let tracker = GLOBAL_TRACKER
.as_ref()
.expect("global tracked marked as initialized, but failed to unwrap");
Some(tracker)
}
}