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