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}