countme/
lib.rs

1//! A library to quickly get the live/total/max counts of allocated instances.
2//!
3//! # Example
4//!
5//! ```
6//! # if cfg!(not(feature = "enable")) { return; }
7//!
8//! #[derive(Default)]
9//! struct Widget {
10//!   _c: countme::Count<Self>,
11//! }
12//!
13//! countme::enable(true);
14//!
15//! let w1 = Widget::default();
16//! let w2 = Widget::default();
17//! let w3 = Widget::default();
18//! drop(w1);
19//!
20//! let counts = countme::get::<Widget>();
21//! assert_eq!(counts.live, 2);
22//! assert_eq!(counts.max_live, 3);
23//! assert_eq!(counts.total, 3);
24//!
25//! eprintln!("{}", countme::get_all());
26//! ```
27//!
28//! # Configuration
29//!
30//! By default, the implementation compiles to no-ops. Therefore it is possible
31//! to include `Count` fields into library types.
32//!
33//! The `enable` cargo feature ungates the counting code. The feature can be
34//! enabled anywhere in the crate graph.
35//!
36//! At run-time, the counters are controlled with [`enable`] function. Counting
37//! is enabled by default if `print_at_exit` feature is enabled. Otherwise
38//! counting is disabled by default. Call `enable(true)` early in `main` to enable:
39//!
40//! ```rust
41//! fn main() {
42//!     countme::enable(std::env::var("COUNTME").is_ok());
43//! }
44//! ```
45//!
46//! The code is optimized for the case where counting is not enabled at runtime
47//! (counting is a relaxed load and a branch to a function call).
48//!
49//! The `print_at_exit` Cargo feature uses `atexit` call to print final counts
50//! before the program exits (it also enables counting at runtime). Use it only
51//! when you can't modify the main to print counts -- `atexit` is not guaranteed
52//! to work with rust's runtime.
53#[cfg(feature = "enable")]
54mod imp;
55
56use std::{fmt, marker::PhantomData};
57
58#[derive(Debug, Clone, PartialEq, Eq, Default)]
59#[non_exhaustive]
60pub struct Counts {
61    /// The total number of tokens created.
62    pub total: usize,
63    /// The historical maximum of the `live` count.
64    pub max_live: usize,
65    /// The number of tokens which were created, but are not destroyed yet.
66    pub live: usize,
67}
68
69/// Store this inside your struct as `_c: countme::Count<Self>`.
70#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
71pub struct Count<T: 'static> {
72    ghost: PhantomData<fn(T)>,
73}
74
75impl<T: 'static> Default for Count<T> {
76    #[inline]
77    fn default() -> Self {
78        Self::new()
79    }
80}
81
82impl<T: 'static> Clone for Count<T> {
83    #[inline]
84    fn clone(&self) -> Self {
85        Self::new()
86    }
87}
88
89impl<T: 'static> Count<T> {
90    /// Create new `Count`, incrementing the corresponding count.
91    #[inline]
92    pub fn new() -> Count<T> {
93        #[cfg(feature = "enable")]
94        imp::inc::<T>();
95        Count { ghost: PhantomData }
96    }
97}
98
99impl<T: 'static> Drop for Count<T> {
100    #[inline]
101    fn drop(&mut self) {
102        #[cfg(feature = "enable")]
103        imp::dec::<T>();
104    }
105}
106
107/// Enable or disable counting at runtime.
108///
109/// Counting is enabled by default if `print_at_exit` feature is enabled.
110/// Otherwise counting is disabled by default.
111///
112/// If neither `enable` nor `print_at_exit` features are enabled, then this function is noop.
113pub fn enable(_yes: bool) {
114    #[cfg(feature = "enable")]
115    imp::enable(_yes);
116}
117
118/// Returns the counts for the `T` type.
119#[inline]
120pub fn get<T: 'static>() -> Counts {
121    #[cfg(feature = "enable")]
122    {
123        return imp::get::<T>();
124    }
125    #[cfg(not(feature = "enable"))]
126    {
127        return Counts::default();
128    }
129}
130
131/// Returns a collection of counts for all types.
132pub fn get_all() -> AllCounts {
133    #[cfg(feature = "enable")]
134    {
135        return imp::get_all();
136    }
137    #[cfg(not(feature = "enable"))]
138    {
139        return AllCounts::default();
140    }
141}
142
143/// A collection of counts for all types.
144#[derive(Default, Clone, Debug)]
145pub struct AllCounts {
146    entries: Vec<(&'static str, Counts)>,
147}
148
149impl fmt::Display for AllCounts {
150    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
151        fn sep(mut n: usize) -> String {
152            let mut groups = Vec::new();
153            while n >= 1000 {
154                groups.push(format!("{:03}", n % 1000));
155                n /= 1000;
156            }
157            groups.push(n.to_string());
158            groups.reverse();
159            groups.join("_")
160        }
161
162        if self.entries.is_empty() {
163            return if cfg!(feature = "enable") {
164                writeln!(f, "all counts are zero")
165            } else {
166                writeln!(f, "counts are disabled")
167            };
168        }
169        let max_width =
170            self.entries.iter().map(|(name, _count)| name.chars().count()).max().unwrap_or(0);
171        for (name, counts) in &self.entries {
172            writeln!(
173                f,
174                "{:<max_width$}  {:>12} {:>12} {:>12}",
175                name,
176                sep(counts.total),
177                sep(counts.max_live),
178                sep(counts.live),
179                max_width = max_width
180            )?;
181        }
182        writeln!(
183            f,
184            "{:<max_width$}  {:>12} {:>12} {:>12}",
185            "",
186            "total",
187            "max_live",
188            "live",
189            max_width = max_width
190        )?;
191        Ok(())
192    }
193}