diff options
Diffstat (limited to 'vendor/countme/src')
| -rw-r--r-- | vendor/countme/src/imp.rs | 175 | ||||
| -rw-r--r-- | vendor/countme/src/lib.rs | 193 |
2 files changed, 368 insertions, 0 deletions
diff --git a/vendor/countme/src/imp.rs b/vendor/countme/src/imp.rs new file mode 100644 index 000000000..c1ace0da7 --- /dev/null +++ b/vendor/countme/src/imp.rs @@ -0,0 +1,175 @@ +use std::{ + any::{type_name, TypeId}, + cell::RefCell, + collections::HashMap, + hash::BuildHasherDefault, + os::raw::c_int, + sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed}, + sync::Arc, +}; + +use dashmap::DashMap; +use once_cell::sync::OnceCell; +use rustc_hash::FxHasher; + +use crate::{AllCounts, Counts}; + +static ENABLE: AtomicBool = AtomicBool::new(cfg!(feature = "print_at_exit")); + +type GlobalStore = DashMap<TypeId, Arc<Store>, BuildHasherDefault<FxHasher>>; + +#[inline] +fn global_store() -> &'static GlobalStore { + static MAP: OnceCell<GlobalStore> = OnceCell::new(); + MAP.get_or_init(|| { + if cfg!(feature = "print_at_exit") { + extern "C" { + fn atexit(f: extern "C" fn()) -> c_int; + } + extern "C" fn print_at_exit() { + eprint!("{}", get_all()); + } + unsafe { + atexit(print_at_exit); + } + } + + GlobalStore::default() + }) +} + +thread_local! { + static LOCAL: RefCell<HashMap<TypeId, Arc<Store>, BuildHasherDefault<FxHasher>>> = RefCell::default(); +} + +pub(crate) fn enable(yes: bool) { + ENABLE.store(yes, Relaxed); +} + +#[inline] +fn enabled() -> bool { + ENABLE.load(Relaxed) +} + +#[inline] +pub(crate) fn dec<T: 'static>() { + if enabled() { + do_dec(TypeId::of::<T>()) + } +} +#[inline(never)] +fn do_dec(key: TypeId) { + LOCAL.with(|local| { + // Fast path: we have needed store in thread local map + if let Some(store) = local.borrow().get(&key) { + store.dec(); + return; + } + + let global = global_store(); + + // Slightly slower: we don't have needed store in our thread local map, + // but some other thread has already initialized the needed store in the global map + if let Some(store) = global.get(&key) { + let store = store.value(); + local.borrow_mut().insert(key, Arc::clone(store)); + store.inc(); + return; + } + + // We only decrement counter after incremenrting it, so this line is unreachable + }) +} + +#[inline] +pub(crate) fn inc<T: 'static>() { + if enabled() { + do_inc(TypeId::of::<T>(), type_name::<T>()) + } +} +#[inline(never)] +fn do_inc(key: TypeId, name: &'static str) { + LOCAL.with(|local| { + // Fast path: we have needed store in thread local map + if let Some(store) = local.borrow().get(&key) { + store.inc(); + return; + } + + let global = global_store(); + + let copy = match global.get(&key) { + // Slightly slower path: we don't have needed store in our thread local map, + // but some other thread has already initialized the needed store in the global map + Some(store) => { + let store = store.value(); + store.inc(); + Arc::clone(store) + } + // Slow path: we are the first to initialize both global and local maps + None => { + let store = global + .entry(key) + .or_insert_with(|| Arc::new(Store { name, ..Store::default() })) + .downgrade(); + let store = store.value(); + + store.inc(); + Arc::clone(store) + } + }; + + local.borrow_mut().insert(key, copy); + }); +} + +pub(crate) fn get<T: 'static>() -> Counts { + do_get(TypeId::of::<T>()) +} +fn do_get(key: TypeId) -> Counts { + global_store().entry(key).or_default().value().read() +} + +pub(crate) fn get_all() -> AllCounts { + let mut entries = global_store() + .iter() + .map(|entry| { + let store = entry.value(); + (store.type_name(), store.read()) + }) + .collect::<Vec<_>>(); + entries.sort_by_key(|(name, _counts)| *name); + AllCounts { entries } +} + +#[derive(Default)] +struct Store { + total: AtomicUsize, + max_live: AtomicUsize, + live: AtomicUsize, + name: &'static str, +} + +impl Store { + fn inc(&self) { + self.total.fetch_add(1, Relaxed); + let live = self.live.fetch_add(1, Relaxed) + 1; + self.max_live.fetch_max(live, Relaxed); + } + + fn dec(&self) { + self.live.fetch_sub(1, Relaxed); + } + + fn read(&self) -> Counts { + Counts { + total: self.total.load(Relaxed), + max_live: self.max_live.load(Relaxed), + live: self.live.load(Relaxed), + } + } + + fn type_name(&self) -> &'static str { + self.name + } +} diff --git a/vendor/countme/src/lib.rs b/vendor/countme/src/lib.rs new file mode 100644 index 000000000..f89e63c5b --- /dev/null +++ b/vendor/countme/src/lib.rs @@ -0,0 +1,193 @@ +//! A library to quickly get the live/total/max counts of allocated instances. +//! +//! # Example +//! +//! ``` +//! # if cfg!(not(feature = "enable")) { return; } +//! +//! #[derive(Default)] +//! struct Widget { +//! _c: countme::Count<Self>, +//! } +//! +//! countme::enable(true); +//! +//! let w1 = Widget::default(); +//! let w2 = Widget::default(); +//! let w3 = Widget::default(); +//! drop(w1); +//! +//! let counts = countme::get::<Widget>(); +//! assert_eq!(counts.live, 2); +//! assert_eq!(counts.max_live, 3); +//! assert_eq!(counts.total, 3); +//! +//! eprintln!("{}", countme::get_all()); +//! ``` +//! +//! # Configuration +//! +//! By default, the implementation compiles to no-ops. Therefore it is possible +//! to include `Count` fields into library types. +//! +//! The `enable` cargo feature ungates the counting code. The feature can be +//! enabled anywhere in the crate graph. +//! +//! At run-time, the counters are controlled with [`enable`] function. Counting +//! is enabled by default if `print_at_exit` feature is enabled. Otherwise +//! counting is disabled by default. Call `enable(true)` early in `main` to enable: +//! +//! ```rust +//! fn main() { +//! countme::enable(std::env::var("COUNTME").is_ok()); +//! } +//! ``` +//! +//! The code is optimized for the case where counting is not enabled at runtime +//! (counting is a relaxed load and a branch to a function call). +//! +//! The `print_at_exit` Cargo feature uses `atexit` call to print final counts +//! before the program exits (it also enables counting at runtime). Use it only +//! when you can't modify the main to print counts -- `atexit` is not guaranteed +//! to work with rust's runtime. +#[cfg(feature = "enable")] +mod imp; + +use std::{fmt, marker::PhantomData}; + +#[derive(Debug, Clone, PartialEq, Eq, Default)] +#[non_exhaustive] +pub struct Counts { + /// The total number of tokens created. + pub total: usize, + /// The historical maximum of the `live` count. + pub max_live: usize, + /// The number of tokens which were created, but are not destroyed yet. + pub live: usize, +} + +/// Store this inside your struct as `_c: countme::Count<Self>`. +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct Count<T: 'static> { + ghost: PhantomData<fn(T)>, +} + +impl<T: 'static> Default for Count<T> { + #[inline] + fn default() -> Self { + Self::new() + } +} + +impl<T: 'static> Clone for Count<T> { + #[inline] + fn clone(&self) -> Self { + Self::new() + } +} + +impl<T: 'static> Count<T> { + /// Create new `Count`, incrementing the corresponding count. + #[inline] + pub fn new() -> Count<T> { + #[cfg(feature = "enable")] + imp::inc::<T>(); + Count { ghost: PhantomData } + } +} + +impl<T: 'static> Drop for Count<T> { + #[inline] + fn drop(&mut self) { + #[cfg(feature = "enable")] + imp::dec::<T>(); + } +} + +/// Enable or disable counting at runtime. +/// +/// Counting is enabled by default if `print_at_exit` feature is enabled. +/// Otherwise counting is disabled by default. +/// +/// If neither `enable` nor `print_at_exit` features are enabled, then this function is noop. +pub fn enable(_yes: bool) { + #[cfg(feature = "enable")] + imp::enable(_yes); +} + +/// Returns the counts for the `T` type. +#[inline] +pub fn get<T: 'static>() -> Counts { + #[cfg(feature = "enable")] + { + return imp::get::<T>(); + } + #[cfg(not(feature = "enable"))] + { + return Counts::default(); + } +} + +/// Returns a collection of counts for all types. +pub fn get_all() -> AllCounts { + #[cfg(feature = "enable")] + { + return imp::get_all(); + } + #[cfg(not(feature = "enable"))] + { + return AllCounts::default(); + } +} + +/// A collection of counts for all types. +#[derive(Default, Clone, Debug)] +pub struct AllCounts { + entries: Vec<(&'static str, Counts)>, +} + +impl fmt::Display for AllCounts { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + fn sep(mut n: usize) -> String { + let mut groups = Vec::new(); + while n >= 1000 { + groups.push(format!("{:03}", n % 1000)); + n /= 1000; + } + groups.push(n.to_string()); + groups.reverse(); + groups.join("_") + } + + if self.entries.is_empty() { + return if cfg!(feature = "enable") { + writeln!(f, "all counts are zero") + } else { + writeln!(f, "counts are disabled") + }; + } + let max_width = + self.entries.iter().map(|(name, _count)| name.chars().count()).max().unwrap_or(0); + for (name, counts) in &self.entries { + writeln!( + f, + "{:<max_width$} {:>12} {:>12} {:>12}", + name, + sep(counts.total), + sep(counts.max_live), + sep(counts.live), + max_width = max_width + )?; + } + writeln!( + f, + "{:<max_width$} {:>12} {:>12} {:>12}", + "", + "total", + "max_live", + "live", + max_width = max_width + )?; + Ok(()) + } +} |