allocation_counter/
lib.rs

1/*!
2This crate provides a method to measure memory allocations while running some code.
3
4It can be used either exploratory (obtaining insights in how much memory allocations
5are being made), or as a tool to assert desired allocation behaviour in tests.
6
7# Usage
8Add as a dependency - since including the trait replaces the global memory allocator,
9you most likely want it gated behind a feature.
10
11```toml
12[features]
13count-allocations = ["allocation-counter"]
14
15[dependencies]
16allocation-counter = { version = "0", optional = true }
17```
18
19The [measure()] function is now available, which can measure memory allocations made
20when the supplied function or closure runs.
21
22Tests can be conditional on the feature:
23
24```
25#[cfg(feature = "count-allocations")]
26#[test]
27{
28    // [...]
29}
30```
31
32The test code itself could look like:
33
34```no_run
35# fn code_that_should_not_allocate() {}
36# fn code_that_should_allocate_a_little() {}
37# fn external_code_that_should_not_be_tested() {}
38// Verify that no memory allocations are made:
39let info = allocation_counter::measure(|| {
40    code_that_should_not_allocate();
41});
42assert_eq!(info.count_total, 0);
43
44// Let's use a case where some allocations are expected.
45let info = allocation_counter::measure(|| {
46    code_that_should_allocate_a_little();
47});
48
49// Using a lower bound can help track behaviour over time:
50assert!((500..600).contains(&info.count_total));
51assert!((10_000..20_000).contains(&info.bytes_total));
52
53// Limit peak memory usage:
54assert!((100..200).contains(&info.count_max));
55assert!((1_000..2_000).contains(&info.bytes_max));
56
57// We don't want any leaks:
58assert_eq!(0, info.count_current);
59assert_eq!(0, info.bytes_current);
60
61// It's possible to opt out of counting allocations
62// for certain parts of the code flow:
63let info = allocation_counter::measure(|| {
64    code_that_should_not_allocate();
65    allocation_counter::opt_out(|| {
66        external_code_that_should_not_be_tested();
67    });
68    code_that_should_not_allocate();
69});
70assert_eq!(0, info.count_total);
71```
72
73Run the tests with the necessary feature enabled.
74
75```sh
76cargo test --features count-allocations
77```
78*/
79
80pub(crate) mod allocator;
81
82/// The allocation information obtained by a [measure()] call.
83#[derive(Clone, Copy, Default, Debug, PartialEq, Eq, Hash)]
84pub struct AllocationInfo {
85    /// The total number of allocations made during a [measure()] call.
86    pub count_total: u64,
87
88    /// The current (net result) number of allocations during a [measure()] call.
89    ///
90    /// A non-zero value of this field means that the function did not deallocate all allocations, as shown below.
91    ///
92    /// ```
93    /// let info = allocation_counter::measure(|| {
94    ///     let b = std::hint::black_box(Box::new(1_u32));
95    ///     std::mem::forget(b);
96    /// });
97    /// assert_eq!(info.count_current, 1);
98    /// ```
99    pub count_current: i64,
100
101    /// The max number of allocations held during a point in time during a [measure()] call.
102    pub count_max: u64,
103
104    /// The total amount of bytes allocated during a [measure()] call.
105    pub bytes_total: u64,
106
107    /// The current (net result) amount of bytes allocated during a [measure()] call.
108    ///
109    /// A non-zero value of this field means that not all memory was deallocated, as shown below.
110    ///
111    /// ```
112    /// let info = allocation_counter::measure(|| {
113    ///     let b = std::hint::black_box(Box::new(1_u32));
114    ///     std::mem::forget(b);
115    /// });
116    /// assert_eq!(info.bytes_current, 4);
117    /// ```
118    pub bytes_current: i64,
119
120    /// The max amount of bytes allocated at one time during a [measure()] call.
121    pub bytes_max: u64,
122}
123
124impl std::ops::AddAssign for AllocationInfo {
125    fn add_assign(&mut self, other: Self) {
126        self.count_total += other.count_total;
127        self.count_current += other.count_current;
128        self.count_max += other.count_max;
129        self.bytes_total += other.bytes_total;
130        self.bytes_current += other.bytes_current;
131        self.bytes_max += other.bytes_max;
132    }
133}
134
135/// Run a closure or function while measuring the performed memory allocations.
136///
137/// Will only measure those allocations done by the current thread, so take care
138/// when interpreting the returned count for multithreaded code.
139///
140/// Use [opt_out()] to opt of of counting allocations temporarily.
141///
142/// Nested `measure()` calls are supported up to a max depth of 64.
143///
144/// # Arguments
145///
146/// - `run_while_measuring` - The code to run while measuring allocations
147///
148/// # Examples
149///
150/// ```
151/// # fn code_that_should_not_allocate_memory() {}
152/// let actual = allocation_counter::measure(|| {
153///      "hello, world".to_string();
154/// });
155/// let expected = allocation_counter::AllocationInfo {
156///     count_total: 1,
157///     count_current: 0,
158///     count_max: 1,
159///     bytes_total: 12,
160///     bytes_current: 0,
161///     bytes_max: 12,
162/// };
163/// assert_eq!(actual, expected);
164/// ```
165pub fn measure<F: FnOnce()>(run_while_measuring: F) -> AllocationInfo {
166    allocator::ALLOCATIONS.with(|info_stack| {
167        let mut info_stack = info_stack.borrow_mut();
168        info_stack.depth += 1;
169        assert!(
170            (info_stack.depth as usize) < allocator::MAX_DEPTH,
171            "Too deep allocation measuring nesting"
172        );
173        let depth = info_stack.depth;
174        info_stack.elements[depth as usize] = AllocationInfo::default();
175    });
176
177    run_while_measuring();
178
179    allocator::ALLOCATIONS.with(|info_stack| {
180        let mut info_stack = info_stack.borrow_mut();
181        let depth = info_stack.depth;
182        let popped = info_stack.elements[depth as usize];
183        info_stack.depth -= 1;
184        let depth = info_stack.depth as usize;
185        info_stack.elements[depth] += popped;
186        popped
187    })
188}
189
190/// Opt out of counting allocations while running some code.
191///
192/// Useful to avoid certain parts of the code flow that should not be counted.
193///
194/// # Arguments
195///
196/// - `run_while_not_counting` - The code to run while not counting allocations
197///
198/// # Examples
199///
200/// ```
201/// # fn code_that_should_not_allocate() {}
202/// # fn external_code_that_should_not_be_tested() {}
203/// let info = allocation_counter::measure(|| {
204///     code_that_should_not_allocate();
205///     allocation_counter::opt_out(|| {
206///         external_code_that_should_not_be_tested();
207///     });
208///     code_that_should_not_allocate();
209/// });
210/// assert_eq!(info.count_total, 0);
211/// ```
212pub fn opt_out<F: FnOnce()>(run_while_not_counting: F) {
213    allocator::DO_COUNT.with(|b| {
214        *b.borrow_mut() += 1;
215        run_while_not_counting();
216        *b.borrow_mut() -= 1;
217    });
218}
219
220#[test]
221fn test_measure() {
222    let info = measure(|| {
223        // Do nothing.
224    });
225    assert_eq!(info.bytes_current, 0);
226    assert_eq!(info.bytes_total, 0);
227    assert_eq!(info.bytes_max, 0);
228    assert_eq!(info.count_current, 0);
229    assert_eq!(info.count_total, 0);
230    assert_eq!(info.count_max, 0);
231
232    let info = measure(|| {
233        {
234            let _a = std::hint::black_box(Box::new(1_u32));
235        }
236        {
237            let _b = std::hint::black_box(Box::new(1_u32));
238        }
239    });
240    assert_eq!(info.bytes_current, 0);
241    assert_eq!(info.bytes_total, 8);
242    assert_eq!(info.bytes_max, 4);
243    assert_eq!(info.count_current, 0);
244    assert_eq!(info.count_total, 2);
245    assert_eq!(info.count_max, 1);
246
247    let info = measure(|| {
248        {
249            let _a = std::hint::black_box(Box::new(1_u32));
250        }
251        let b = std::hint::black_box(Box::new(1_u32));
252        std::mem::forget(b);
253    });
254    assert_eq!(info.bytes_current, 4);
255    assert_eq!(info.bytes_total, 8);
256    assert_eq!(info.bytes_max, 4);
257    assert_eq!(info.count_current, 1);
258    assert_eq!(info.count_total, 2);
259    assert_eq!(info.count_max, 1);
260
261    let info = measure(|| {
262        let a = std::hint::black_box(Box::new(1_u32));
263        let b = std::hint::black_box(Box::new(1_u32));
264        let _c = std::hint::black_box(Box::new(*a + *b));
265    });
266    assert_eq!(info.bytes_current, 0);
267    assert_eq!(info.bytes_total, 12);
268    assert_eq!(info.bytes_max, 12);
269    assert_eq!(info.count_current, 0);
270    assert_eq!(info.count_total, 3);
271    assert_eq!(info.count_max, 3);
272}
273
274#[test]
275fn test_opt_out() {
276    let allocations = measure(|| {
277        // Do nothing.
278    });
279    assert_eq!(allocations.count_total, 0);
280
281    let allocations = measure(|| {
282        let v: Vec<u32> = vec![12];
283        assert_eq!(v.len(), 1);
284        opt_out(|| {
285            let v: Vec<u32> = vec![12];
286            assert_eq!(v.len(), 1);
287            opt_out(|| {
288                let v: Vec<u32> = vec![12];
289                assert_eq!(v.len(), 1);
290            });
291        });
292        let v: Vec<u32> = vec![12];
293        assert_eq!(v.len(), 1);
294        let v: Vec<u32> = vec![12];
295        assert_eq!(v.len(), 1);
296    });
297    assert_eq!(allocations.count_total, 3);
298
299    let info = measure(|| {
300        opt_out(|| {
301            let v: Vec<u32> = vec![12];
302            assert_eq!(v.len(), 1);
303        });
304    });
305    assert_eq!(0, info.count_total);
306}
307
308#[test]
309fn test_nested_counting() {
310    let info = measure(|| {
311        let _a = std::hint::black_box(Box::new(1_u32));
312        let info = measure(|| {
313            let _b = std::hint::black_box(Box::new(1_u32));
314        });
315        assert_eq!(info.bytes_current, 0);
316        assert_eq!(info.bytes_total, 4);
317        assert_eq!(info.bytes_max, 4);
318        assert_eq!(info.count_current, 0);
319        assert_eq!(info.count_total, 1);
320        assert_eq!(info.count_max, 1);
321    });
322    assert_eq!(info.bytes_current, 0);
323    assert_eq!(info.bytes_total, 8);
324    assert_eq!(info.bytes_max, 8);
325    assert_eq!(info.count_current, 0);
326    assert_eq!(info.count_total, 2);
327    assert_eq!(info.count_max, 2);
328}