mod_alloc/lib.rs
1//! # mod-alloc
2//!
3//! Allocation profiling for Rust. Tracks allocation counts, total
4//! bytes, peak resident memory, and current resident memory by
5//! wrapping the system allocator via [`GlobalAlloc`].
6//!
7//! Designed as a lean replacement for `dhat` with MSRV 1.75 and
8//! zero external dependencies on the hot path.
9//!
10//! ## Installing as the global allocator
11//!
12//! ```no_run
13//! use mod_alloc::{ModAlloc, Profiler};
14//!
15//! #[global_allocator]
16//! static GLOBAL: ModAlloc = ModAlloc::new();
17//!
18//! fn main() {
19//! let p = Profiler::start();
20//!
21//! let v: Vec<u64> = (0..1000).collect();
22//! drop(v);
23//!
24//! let stats = p.stop();
25//! println!("Allocations: {}", stats.alloc_count);
26//! println!("Total bytes: {}", stats.total_bytes);
27//! println!("Peak bytes (absolute): {}", stats.peak_bytes);
28//! }
29//! ```
30//!
31//! ## Counter semantics
32//!
33//! The four Tier 1 counters track allocator activity since the
34//! installed [`ModAlloc`] began counting (or since the last
35//! [`ModAlloc::reset`] call):
36//!
37//! | Counter | Updated on `alloc` | Updated on `dealloc` |
38//! |-----------------|-------------------------------|----------------------|
39//! | `alloc_count` | `+= 1` | (unchanged) |
40//! | `total_bytes` | `+= size` | (unchanged) |
41//! | `current_bytes` | `+= size` | `-= size` |
42//! | `peak_bytes` | high-water mark of `current` | (unchanged) |
43//!
44//! `realloc` is counted as one allocation event. `total_bytes`
45//! increases by the growth delta on a growing realloc and is
46//! unchanged on a shrinking realloc.
47//!
48//! ## Status
49//!
50//! v0.9.1 added Tier 2 (inline backtrace capture) behind the
51//! `backtraces` feature. v0.9.2 added symbolication behind the
52//! `symbolicate` feature. v0.9.3 wired up `dhat-compat` to emit
53//! the per-call-site report as DHAT-format JSON that the upstream
54//! `dh_view.html` viewer loads directly. v0.9.4 added the
55//! `dhat_compat` drop-in surface for migrating from `dhat-rs`.
56//! v0.9.5 cut Tier 2 per-allocation overhead by ~32x by removing
57//! the per-thread arena layer. Default builds still ship Tier 1
58//! counters only.
59//!
60//! ## Backtraces (`backtraces` feature)
61//!
62//! With `mod-alloc = { version = "0.9", features = ["backtraces"] }`
63//! and `RUSTFLAGS="-C force-frame-pointers=yes"`, each tracked
64//! allocation captures up to 8 frames of its call site via inline
65//! frame-pointer walking on `x86_64` and `aarch64`. Per-call-site
66//! aggregation is exposed via `ModAlloc::call_sites` (available
67//! only with the `backtraces` feature); the result is raw return
68//! addresses. Symbolication ships in v0.9.2.
69//!
70//! Aggregation-table size is controlled by the `MOD_ALLOC_BUCKETS`
71//! environment variable at process start (default 4,096 buckets,
72//! ~384 KB).
73
74#![cfg_attr(docsrs, feature(doc_cfg))]
75#![warn(missing_docs)]
76#![warn(rust_2018_idioms)]
77
78#[cfg(feature = "backtraces")]
79mod backtrace;
80
81#[cfg(feature = "backtraces")]
82pub use backtrace::CallSiteStats;
83
84#[cfg(feature = "symbolicate")]
85mod symbolicate;
86
87#[cfg(feature = "symbolicate")]
88pub use symbolicate::{SymbolicatedCallSite, SymbolicatedFrame};
89
90#[cfg(feature = "dhat-compat")]
91mod dhat_json;
92
93#[cfg(feature = "dhat-compat")]
94pub mod dhat_compat;
95
96use std::alloc::{GlobalAlloc, Layout, System};
97use std::cell::Cell;
98use std::ptr;
99use std::sync::atomic::{AtomicPtr, AtomicU64, Ordering};
100
101// Process-wide handle to the installed `ModAlloc`. Populated lazily
102// on the first non-reentrant alloc call. `Profiler` reads from this
103// to locate the canonical counters without requiring an explicit
104// registration call from the user.
105static GLOBAL_HANDLE: AtomicPtr<ModAlloc> = AtomicPtr::new(ptr::null_mut());
106
107thread_local! {
108 // Reentrancy flag. Set while inside the tracking path; if any
109 // allocation occurs while set, the recursive call bypasses
110 // tracking and forwards directly to the System allocator.
111 //
112 // `const` initialization (stable since 1.59) avoids any lazy
113 // construction allocation inside the TLS access path.
114 static IN_ALLOC: Cell<bool> = const { Cell::new(false) };
115}
116
117// RAII guard for the reentrancy flag. `enter` returns `None` if the
118// current thread is already inside a tracked allocation (caller
119// must skip counter updates) or if TLS is unavailable (e.g. during
120// thread teardown). The guard clears the flag on drop.
121struct ReentryGuard;
122
123impl ReentryGuard {
124 fn enter() -> Option<Self> {
125 IN_ALLOC
126 .try_with(|flag| {
127 if flag.get() {
128 None
129 } else {
130 flag.set(true);
131 Some(ReentryGuard)
132 }
133 })
134 .ok()
135 .flatten()
136 }
137}
138
139impl Drop for ReentryGuard {
140 fn drop(&mut self) {
141 let _ = IN_ALLOC.try_with(|flag| flag.set(false));
142 }
143}
144
145/// Global allocator wrapper that tracks allocations.
146///
147/// Install as `#[global_allocator]` to enable tracking. The wrapper
148/// forwards every allocation, deallocation, reallocation, and
149/// zero-initialised allocation to [`std::alloc::System`] and records
150/// the event in four lock-free [`AtomicU64`] counters.
151///
152/// # Example
153///
154/// ```no_run
155/// use mod_alloc::ModAlloc;
156///
157/// #[global_allocator]
158/// static GLOBAL: ModAlloc = ModAlloc::new();
159///
160/// fn main() {
161/// let v: Vec<u8> = vec![0; 1024];
162/// let stats = GLOBAL.snapshot();
163/// assert!(stats.alloc_count >= 1);
164/// drop(v);
165/// }
166/// ```
167pub struct ModAlloc {
168 alloc_count: AtomicU64,
169 total_bytes: AtomicU64,
170 peak_bytes: AtomicU64,
171 current_bytes: AtomicU64,
172 live_count: AtomicU64,
173 peak_live_count: AtomicU64,
174}
175
176impl ModAlloc {
177 /// Construct a new `ModAlloc` allocator wrapper.
178 ///
179 /// All counters start at zero. This function is `const`, which
180 /// allows construction in a `static` for use as
181 /// `#[global_allocator]`.
182 ///
183 /// # Example
184 ///
185 /// ```
186 /// use mod_alloc::ModAlloc;
187 ///
188 /// static GLOBAL: ModAlloc = ModAlloc::new();
189 /// let stats = GLOBAL.snapshot();
190 /// assert_eq!(stats.alloc_count, 0);
191 /// ```
192 pub const fn new() -> Self {
193 Self {
194 alloc_count: AtomicU64::new(0),
195 total_bytes: AtomicU64::new(0),
196 peak_bytes: AtomicU64::new(0),
197 current_bytes: AtomicU64::new(0),
198 live_count: AtomicU64::new(0),
199 peak_live_count: AtomicU64::new(0),
200 }
201 }
202
203 /// Snapshot the current counter values.
204 ///
205 /// Each counter is read independently with `Relaxed` ordering;
206 /// the resulting [`AllocStats`] is a coherent best-effort view
207 /// but does not represent a single atomic moment in time. For
208 /// scoped measurement, prefer [`Profiler`].
209 ///
210 /// # Example
211 ///
212 /// ```
213 /// use mod_alloc::ModAlloc;
214 ///
215 /// let alloc = ModAlloc::new();
216 /// let stats = alloc.snapshot();
217 /// assert_eq!(stats.alloc_count, 0);
218 /// ```
219 pub fn snapshot(&self) -> AllocStats {
220 AllocStats {
221 alloc_count: self.alloc_count.load(Ordering::Relaxed),
222 total_bytes: self.total_bytes.load(Ordering::Relaxed),
223 peak_bytes: self.peak_bytes.load(Ordering::Relaxed),
224 current_bytes: self.current_bytes.load(Ordering::Relaxed),
225 live_count: self.live_count.load(Ordering::Relaxed),
226 peak_live_count: self.peak_live_count.load(Ordering::Relaxed),
227 }
228 }
229
230 /// Reset all counters to zero.
231 ///
232 /// Intended for use at the start of a profile run, before any
233 /// outstanding allocations exist. Calling `reset` while
234 /// allocations are live can cause `current_bytes` to wrap on
235 /// subsequent deallocations; the other counters are unaffected.
236 ///
237 /// # Example
238 ///
239 /// ```
240 /// use mod_alloc::ModAlloc;
241 ///
242 /// let alloc = ModAlloc::new();
243 /// alloc.reset();
244 /// let stats = alloc.snapshot();
245 /// assert_eq!(stats.alloc_count, 0);
246 /// ```
247 pub fn reset(&self) {
248 self.alloc_count.store(0, Ordering::Relaxed);
249 self.total_bytes.store(0, Ordering::Relaxed);
250 self.peak_bytes.store(0, Ordering::Relaxed);
251 self.current_bytes.store(0, Ordering::Relaxed);
252 self.live_count.store(0, Ordering::Relaxed);
253 self.peak_live_count.store(0, Ordering::Relaxed);
254 }
255
256 #[inline]
257 fn record_alloc(&self, size: u64) {
258 self.alloc_count.fetch_add(1, Ordering::Relaxed);
259 self.total_bytes.fetch_add(size, Ordering::Relaxed);
260 let new_current = self.current_bytes.fetch_add(size, Ordering::Relaxed) + size;
261 self.peak_bytes.fetch_max(new_current, Ordering::Relaxed);
262 let new_live = self.live_count.fetch_add(1, Ordering::Relaxed) + 1;
263 self.peak_live_count.fetch_max(new_live, Ordering::Relaxed);
264 }
265
266 #[inline]
267 fn record_dealloc(&self, size: u64) {
268 self.current_bytes.fetch_sub(size, Ordering::Relaxed);
269 self.live_count.fetch_sub(1, Ordering::Relaxed);
270 }
271
272 #[inline]
273 fn record_realloc(&self, old_size: u64, new_size: u64) {
274 self.alloc_count.fetch_add(1, Ordering::Relaxed);
275 if new_size > old_size {
276 let delta = new_size - old_size;
277 self.total_bytes.fetch_add(delta, Ordering::Relaxed);
278 let new_current = self.current_bytes.fetch_add(delta, Ordering::Relaxed) + delta;
279 self.peak_bytes.fetch_max(new_current, Ordering::Relaxed);
280 } else if new_size < old_size {
281 self.current_bytes
282 .fetch_sub(old_size - new_size, Ordering::Relaxed);
283 }
284 }
285
286 #[inline]
287 fn register_self(&self) {
288 if GLOBAL_HANDLE.load(Ordering::Relaxed).is_null() {
289 let _ = GLOBAL_HANDLE.compare_exchange(
290 ptr::null_mut(),
291 self as *const ModAlloc as *mut ModAlloc,
292 Ordering::Release,
293 Ordering::Relaxed,
294 );
295 }
296 }
297
298 /// Drain the per-call-site aggregation table into a `Vec`.
299 ///
300 /// Available only with the `backtraces` cargo feature. The
301 /// returned vector contains one [`CallSiteStats`] per unique
302 /// call site observed since the table was first written. Each
303 /// row carries up to 8 raw return addresses (top of stack
304 /// first), the number of allocations attributed to that site,
305 /// and the total bytes.
306 ///
307 /// Symbolication (resolving addresses to function names)
308 /// lands in `v0.9.2`. This method exposes raw addresses only.
309 ///
310 /// # Example
311 ///
312 /// ```no_run
313 /// # #[cfg(feature = "backtraces")]
314 /// # fn demo() {
315 /// use mod_alloc::ModAlloc;
316 ///
317 /// #[global_allocator]
318 /// static GLOBAL: ModAlloc = ModAlloc::new();
319 ///
320 /// let _v: Vec<u8> = vec![0; 1024];
321 /// for site in GLOBAL.call_sites() {
322 /// println!("{} allocs, {} bytes at {:#x}",
323 /// site.count, site.total_bytes, site.frames[0]);
324 /// }
325 /// # }
326 /// ```
327 #[cfg(feature = "backtraces")]
328 pub fn call_sites(&self) -> Vec<CallSiteStats> {
329 backtrace::call_sites_report()
330 }
331
332 /// Drain the per-call-site table and symbolicate each frame
333 /// against the running binary's own debug info.
334 ///
335 /// Available only with the `symbolicate` cargo feature, which
336 /// also implies `backtraces`. Returns one
337 /// [`SymbolicatedCallSite`] per unique call site, each
338 /// carrying resolved function names plus (where available)
339 /// source file and line.
340 ///
341 /// Allocates. Safe to call from non-allocator contexts only
342 /// (ordinary user code outside the global-allocator hook).
343 ///
344 /// Results are cached per-address across calls.
345 ///
346 /// # Example
347 ///
348 /// ```no_run
349 /// # #[cfg(feature = "symbolicate")]
350 /// # fn demo() {
351 /// use mod_alloc::ModAlloc;
352 ///
353 /// #[global_allocator]
354 /// static GLOBAL: ModAlloc = ModAlloc::new();
355 ///
356 /// let _v: Vec<u8> = vec![0; 1024];
357 /// for site in GLOBAL.symbolicated_report() {
358 /// let top = &site.frames[0];
359 /// println!("{} allocs / {} bytes at {}",
360 /// site.count,
361 /// site.total_bytes,
362 /// top.function.as_deref().unwrap_or("<unresolved>"));
363 /// }
364 /// # }
365 /// ```
366 #[cfg(feature = "symbolicate")]
367 pub fn symbolicated_report(&self) -> Vec<SymbolicatedCallSite> {
368 symbolicate::symbolicated_report()
369 }
370
371 /// Render the per-call-site report as a DHAT-compatible JSON
372 /// string.
373 ///
374 /// Available only with the `dhat-compat` cargo feature (which
375 /// implies `backtraces`). When the `symbolicate` feature is
376 /// also active, frame strings carry function names and
377 /// (where available) source file and line; otherwise they
378 /// carry raw hex addresses.
379 ///
380 /// Allocates. Safe to call from non-allocator contexts only
381 /// (ordinary user code outside the global-allocator hook).
382 ///
383 /// The output schema (`dhatFileVersion: 2`, `mode: "rust-heap"`)
384 /// matches the format consumed by the upstream
385 /// `dh_view.html` viewer shipped with Valgrind.
386 ///
387 /// # Example
388 ///
389 /// ```no_run
390 /// # #[cfg(feature = "dhat-compat")]
391 /// # fn demo() {
392 /// use mod_alloc::ModAlloc;
393 ///
394 /// #[global_allocator]
395 /// static GLOBAL: ModAlloc = ModAlloc::new();
396 ///
397 /// let _v: Vec<u8> = vec![0; 1024];
398 /// let json = GLOBAL.dhat_json_string();
399 /// assert!(json.contains("\"dhatFileVersion\":2"));
400 /// # }
401 /// ```
402 #[cfg(feature = "dhat-compat")]
403 pub fn dhat_json_string(&self) -> String {
404 dhat_json::dhat_json_string()
405 }
406
407 /// Render the per-call-site report and write it to `path` as
408 /// DHAT-compatible JSON.
409 ///
410 /// Available only with the `dhat-compat` cargo feature.
411 /// Mirrors `dhat-rs`'s convention of writing
412 /// `dhat-heap.json` to the current directory; pass that path
413 /// to drop a file the upstream viewer will load directly.
414 ///
415 /// Allocates. Safe to call from non-allocator contexts only.
416 ///
417 /// # Example
418 ///
419 /// ```no_run
420 /// # #[cfg(feature = "dhat-compat")]
421 /// # fn demo() -> std::io::Result<()> {
422 /// use mod_alloc::ModAlloc;
423 ///
424 /// #[global_allocator]
425 /// static GLOBAL: ModAlloc = ModAlloc::new();
426 ///
427 /// let _v: Vec<u8> = vec![0; 1024];
428 /// GLOBAL.write_dhat_json("dhat-heap.json")?;
429 /// # Ok(())
430 /// # }
431 /// ```
432 #[cfg(feature = "dhat-compat")]
433 pub fn write_dhat_json<P: AsRef<std::path::Path>>(&self, path: P) -> std::io::Result<()> {
434 dhat_json::write_dhat_json(path.as_ref())
435 }
436}
437
438impl Default for ModAlloc {
439 fn default() -> Self {
440 Self::new()
441 }
442}
443
444// SAFETY: `ModAlloc` adds counter bookkeeping but performs all
445// underlying allocation through [`std::alloc::System`]. Each method
446// forwards its arguments unchanged to `System` and only inspects
447// the result; size/alignment invariants required by the
448// `GlobalAlloc` contract are passed through unmodified, so the
449// caller's contract to us becomes our contract to System.
450//
451// The counter-update path uses thread-local reentrancy detection
452// (see `ReentryGuard`) so that any allocation triggered transitively
453// inside the tracking path bypasses tracking and forwards directly
454// to System, preserving the "hook MUST NOT itself allocate"
455// invariant from REPS section 4.
456unsafe impl GlobalAlloc for ModAlloc {
457 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
458 // SAFETY: per `GlobalAlloc::alloc`, `layout` has non-zero
459 // size; we forward unchanged to `System.alloc`, which has
460 // the same contract.
461 let ptr = unsafe { System.alloc(layout) };
462 if !ptr.is_null() {
463 if let Some(_g) = ReentryGuard::enter() {
464 let size = layout.size() as u64;
465 self.record_alloc(size);
466 self.register_self();
467 #[cfg(feature = "backtraces")]
468 backtrace::record_event(size);
469 }
470 }
471 ptr
472 }
473
474 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
475 // SAFETY: same invariants as `alloc`; `layout` forwarded
476 // unchanged. `System.alloc_zeroed` zero-fills the returned
477 // memory, satisfying the `GlobalAlloc::alloc_zeroed`
478 // contract.
479 let ptr = unsafe { System.alloc_zeroed(layout) };
480 if !ptr.is_null() {
481 if let Some(_g) = ReentryGuard::enter() {
482 let size = layout.size() as u64;
483 self.record_alloc(size);
484 self.register_self();
485 #[cfg(feature = "backtraces")]
486 backtrace::record_event(size);
487 }
488 }
489 ptr
490 }
491
492 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
493 // SAFETY: per `GlobalAlloc::dealloc`, `ptr` was returned by a
494 // prior call to `alloc`/`alloc_zeroed`/`realloc` on this
495 // allocator with the given `layout`; we forwarded all of
496 // those to `System` with the same `layout`, so the inverse
497 // pairing for `System.dealloc(ptr, layout)` is valid.
498 unsafe { System.dealloc(ptr, layout) };
499 if let Some(_g) = ReentryGuard::enter() {
500 self.record_dealloc(layout.size() as u64);
501 }
502 }
503
504 unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
505 // SAFETY: per `GlobalAlloc::realloc`, `ptr` was returned by
506 // a prior allocation with `layout`, `new_size` is non-zero,
507 // and the alignment in `layout` remains valid for the new
508 // size. We forward all three to `System.realloc` which has
509 // the same contract.
510 let new_ptr = unsafe { System.realloc(ptr, layout, new_size) };
511 if !new_ptr.is_null() {
512 if let Some(_g) = ReentryGuard::enter() {
513 self.record_realloc(layout.size() as u64, new_size as u64);
514 self.register_self();
515 // Per dhat semantics: realloc records as one event
516 // attributed to `new_size` (including shrinks).
517 #[cfg(feature = "backtraces")]
518 backtrace::record_event(new_size as u64);
519 }
520 }
521 new_ptr
522 }
523}
524
525/// Snapshot of allocation statistics at a point in time.
526///
527/// Produced by [`ModAlloc::snapshot`] and [`Profiler::stop`].
528///
529/// # Example
530///
531/// ```
532/// use mod_alloc::AllocStats;
533///
534/// let stats = AllocStats {
535/// alloc_count: 10,
536/// total_bytes: 1024,
537/// peak_bytes: 512,
538/// current_bytes: 256,
539/// live_count: 4,
540/// peak_live_count: 7,
541/// };
542/// assert_eq!(stats.alloc_count, 10);
543/// ```
544///
545/// # Version note
546///
547/// `live_count` and `peak_live_count` were added in v0.9.4 to
548/// support dhat-rs's `HeapStats` shape via the `dhat-compat`
549/// feature. Callers constructing `AllocStats` via struct literal
550/// must initialise the new fields; callers that only consume
551/// `AllocStats` via [`ModAlloc::snapshot`] or
552/// [`Profiler::stop`] are unaffected.
553#[derive(Debug, Clone, Copy, PartialEq, Eq)]
554pub struct AllocStats {
555 /// Number of allocations performed.
556 pub alloc_count: u64,
557 /// Total bytes allocated across all allocations. Reallocations
558 /// contribute the growth delta (or zero on shrink).
559 pub total_bytes: u64,
560 /// Peak resident bytes (highest `current_bytes` ever observed).
561 pub peak_bytes: u64,
562 /// Currently-allocated bytes (allocations minus deallocations).
563 pub current_bytes: u64,
564 /// Currently-alive allocation count (allocations minus
565 /// deallocations). Mirrors `dhat::HeapStats::curr_blocks`.
566 pub live_count: u64,
567 /// Peak live allocation count (highest `live_count` ever
568 /// observed). Mirrors `dhat::HeapStats::max_blocks`.
569 pub peak_live_count: u64,
570}
571
572/// Scoped profiler that captures a delta between start and stop.
573///
574/// Read the snapshot of the installed [`ModAlloc`] on construction
575/// and again on [`Profiler::stop`], returning the difference. If no
576/// `ModAlloc` is installed as `#[global_allocator]` and no
577/// allocation has occurred through it yet, both snapshots are
578/// zero and the delta is zero.
579///
580/// # Example
581///
582/// ```no_run
583/// use mod_alloc::{ModAlloc, Profiler};
584///
585/// #[global_allocator]
586/// static GLOBAL: ModAlloc = ModAlloc::new();
587///
588/// fn main() {
589/// let p = Profiler::start();
590/// let v: Vec<u8> = vec![0; 1024];
591/// drop(v);
592/// let stats = p.stop();
593/// println!("Captured {} alloc events", stats.alloc_count);
594/// }
595/// ```
596pub struct Profiler {
597 baseline: AllocStats,
598}
599
600impl Profiler {
601 /// Begin profiling, capturing the current allocation state.
602 ///
603 /// If no `ModAlloc` is installed as `#[global_allocator]` or no
604 /// allocation has occurred yet, the captured baseline is all
605 /// zeros.
606 ///
607 /// # Example
608 ///
609 /// ```
610 /// use mod_alloc::Profiler;
611 ///
612 /// let p = Profiler::start();
613 /// let _delta = p.stop();
614 /// ```
615 pub fn start() -> Self {
616 Self {
617 baseline: current_snapshot_or_zeros(),
618 }
619 }
620
621 /// Stop profiling and return the delta from start.
622 ///
623 /// `alloc_count`, `total_bytes`, and `current_bytes` are deltas
624 /// from `start()` to `stop()`. `peak_bytes` is the absolute
625 /// high-water mark observed during the profiling window (peak
626 /// has no meaningful delta semantic).
627 ///
628 /// # Example
629 ///
630 /// ```
631 /// use mod_alloc::Profiler;
632 ///
633 /// let p = Profiler::start();
634 /// let stats = p.stop();
635 /// assert_eq!(stats.alloc_count, 0);
636 /// ```
637 pub fn stop(self) -> AllocStats {
638 let now = current_snapshot_or_zeros();
639 AllocStats {
640 alloc_count: now.alloc_count.saturating_sub(self.baseline.alloc_count),
641 total_bytes: now.total_bytes.saturating_sub(self.baseline.total_bytes),
642 current_bytes: now
643 .current_bytes
644 .saturating_sub(self.baseline.current_bytes),
645 peak_bytes: now.peak_bytes,
646 live_count: now.live_count.saturating_sub(self.baseline.live_count),
647 peak_live_count: now.peak_live_count,
648 }
649 }
650}
651
652pub(crate) fn current_snapshot_or_zeros() -> AllocStats {
653 let p = GLOBAL_HANDLE.load(Ordering::Acquire);
654 if p.is_null() {
655 AllocStats {
656 alloc_count: 0,
657 total_bytes: 0,
658 peak_bytes: 0,
659 current_bytes: 0,
660 live_count: 0,
661 peak_live_count: 0,
662 }
663 } else {
664 // SAFETY: `GLOBAL_HANDLE` is only ever set by
665 // `ModAlloc::register_self` to point at the address of a
666 // `#[global_allocator] static` (or any other `'static`
667 // `ModAlloc`). That target has `'static` lifetime, so the
668 // pointer remains valid for the remainder of the program.
669 // We produce only a shared borrow used to call `&self`
670 // methods that read atomic counters; no mutation through
671 // the pointer occurs here.
672 unsafe { (*p).snapshot() }
673 }
674}
675
676#[cfg(test)]
677mod tests {
678 use super::*;
679
680 #[test]
681 fn allocator_constructs() {
682 let _ = ModAlloc::new();
683 }
684
685 #[test]
686 fn snapshot_returns_zeros_initially() {
687 let a = ModAlloc::new();
688 let s = a.snapshot();
689 assert_eq!(s.alloc_count, 0);
690 assert_eq!(s.total_bytes, 0);
691 assert_eq!(s.peak_bytes, 0);
692 assert_eq!(s.current_bytes, 0);
693 }
694
695 #[test]
696 fn reset_works() {
697 let a = ModAlloc::new();
698 a.reset();
699 let s = a.snapshot();
700 assert_eq!(s.alloc_count, 0);
701 }
702
703 #[test]
704 fn record_alloc_updates_counters() {
705 let a = ModAlloc::new();
706 a.record_alloc(128);
707 a.record_alloc(256);
708 let s = a.snapshot();
709 assert_eq!(s.alloc_count, 2);
710 assert_eq!(s.total_bytes, 384);
711 assert_eq!(s.current_bytes, 384);
712 assert_eq!(s.peak_bytes, 384);
713 }
714
715 #[test]
716 fn record_dealloc_decreases_current_only() {
717 let a = ModAlloc::new();
718 a.record_alloc(1000);
719 a.record_dealloc(400);
720 let s = a.snapshot();
721 assert_eq!(s.alloc_count, 1);
722 assert_eq!(s.total_bytes, 1000);
723 assert_eq!(s.current_bytes, 600);
724 assert_eq!(s.peak_bytes, 1000);
725 // live_count tracks the *count* of alive allocations, not
726 // bytes — a single alloc + a single (partial) dealloc
727 // event leaves zero alive allocations from the counter's
728 // perspective even though current_bytes is non-zero.
729 // This mirrors dhat's curr_blocks semantics where every
730 // alloc/dealloc pair is a 1:1 block lifecycle.
731 assert_eq!(s.live_count, 0);
732 assert_eq!(s.peak_live_count, 1);
733 }
734
735 #[test]
736 fn live_counters_track_alive_blocks() {
737 let a = ModAlloc::new();
738 a.record_alloc(100);
739 a.record_alloc(200);
740 a.record_alloc(300);
741 let s = a.snapshot();
742 assert_eq!(s.live_count, 3);
743 assert_eq!(s.peak_live_count, 3);
744
745 a.record_dealloc(100);
746 a.record_dealloc(200);
747 let s = a.snapshot();
748 assert_eq!(s.live_count, 1);
749 assert_eq!(s.peak_live_count, 3, "peak holds high-water mark");
750 }
751
752 #[test]
753 fn record_realloc_does_not_touch_live_count() {
754 let a = ModAlloc::new();
755 a.record_alloc(100);
756 let pre = a.snapshot();
757 a.record_realloc(100, 250);
758 let post = a.snapshot();
759 assert_eq!(
760 pre.live_count, post.live_count,
761 "realloc keeps the same block, live_count unchanged"
762 );
763 }
764
765 #[test]
766 fn record_realloc_growth_updates_total_and_peak() {
767 let a = ModAlloc::new();
768 a.record_alloc(100);
769 a.record_realloc(100, 250);
770 let s = a.snapshot();
771 assert_eq!(s.alloc_count, 2);
772 assert_eq!(s.total_bytes, 250);
773 assert_eq!(s.current_bytes, 250);
774 assert_eq!(s.peak_bytes, 250);
775 }
776
777 #[test]
778 fn record_realloc_shrink_only_adjusts_current() {
779 let a = ModAlloc::new();
780 a.record_alloc(500);
781 a.record_realloc(500, 200);
782 let s = a.snapshot();
783 assert_eq!(s.alloc_count, 2);
784 assert_eq!(s.total_bytes, 500);
785 assert_eq!(s.current_bytes, 200);
786 assert_eq!(s.peak_bytes, 500);
787 }
788
789 #[test]
790 fn peak_holds_high_water_mark() {
791 let a = ModAlloc::new();
792 a.record_alloc(1000);
793 a.record_dealloc(1000);
794 a.record_alloc(500);
795 let s = a.snapshot();
796 assert_eq!(s.peak_bytes, 1000);
797 assert_eq!(s.current_bytes, 500);
798 }
799
800 #[test]
801 fn reentry_guard_blocks_nested_entry() {
802 let outer = ReentryGuard::enter();
803 assert!(outer.is_some());
804 let inner = ReentryGuard::enter();
805 assert!(inner.is_none(), "nested entry must be denied");
806 drop(outer);
807 let after = ReentryGuard::enter();
808 assert!(after.is_some(), "entry must be allowed after outer drops");
809 }
810
811 #[test]
812 fn profiler_start_stop_with_no_handle() {
813 let p = Profiler::start();
814 let s = p.stop();
815 assert_eq!(s.alloc_count, 0);
816 }
817}