Skip to main content

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}