Skip to main content

mimalloc_pprof/
lib.rs

1//! Rust global allocator support for the in-tree mimalloc build.
2//!
3//! ```no_run
4//! use mimalloc_pprof::{prof, MiMalloc};
5//! #[global_allocator] static ALLOCATOR: MiMalloc = MiMalloc;
6//! # fn main() -> std::io::Result<()> {
7//! prof::start(512 * 1024);
8//! prof::dump_file(std::path::Path::new("heap.prof"))?;
9//! # Ok(()) }
10//! ```
11//!
12//! See the README's Rust integration guide for frame-pointer and line-table
13//! build flags. Open the resulting profile with `pprof -http=: app.exe heap.prof`.
14
15use core::alloc::{GlobalAlloc, Layout};
16use core::ffi::c_void;
17use std::ffi::CString;
18use std::path::PathBuf;
19
20pub mod sys;
21
22/// A `#[global_allocator]` implementation backed by mimalloc.
23pub struct MiMalloc;
24
25unsafe impl GlobalAlloc for MiMalloc {
26    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
27        sys::mi_malloc_aligned(layout.size(), layout.align()).cast()
28    }
29
30    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
31        sys::mi_free(ptr.cast::<c_void>());
32    }
33
34    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
35        sys::mi_realloc_aligned(ptr.cast::<c_void>(), new_size, layout.align()).cast()
36    }
37
38    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
39        sys::mi_zalloc_aligned(layout.size(), layout.align()).cast()
40    }
41}
42
43/// Allocate `size` bytes from mimalloc's raw-OS-layer "unwrapped" path.
44///
45/// Thin wrapper around `mi_unwrapped_malloc` (include/mimalloc/memory-events.h):
46/// backed directly by `_mi_os_alloc_aligned`, never by the hooked `mi_malloc`
47/// family. Page granular, so this is not meant for hot-path/small allocations
48/// — it exists for low-level instrumentation and recursion avoidance (e.g.
49/// scratch storage for a memory-change callback that must not recursively
50/// enter mimalloc). Excluded from normal mimalloc allocation stats and from
51/// the memory-change accounting.
52///
53/// Returns a null pointer on failure (including invalid `alignment`; see
54/// `# Safety` below).
55///
56/// # Safety
57///
58/// - `alignment` must be `0` (treated as `align_of::<*const ()>()`, i.e.
59///   pointer size) or a power of two. A non-power-of-two, non-zero alignment
60///   is a validated input on the C side: `mi_unwrapped_malloc` returns a null
61///   pointer rather than invoking undefined behavior, but callers should not
62///   rely on that as anything other than a defined-failure contract — treat
63///   the alignment argument as a precondition to get right, not a value to
64///   probe.
65/// - The returned pointer, if non-null, must be passed only to
66///   [`unwrapped_free`] or [`unwrapped_realloc`] — never to `mi_free`, this
67///   crate's [`MiMalloc`] allocator, or Rust's global allocator, and vice
68///   versa (a pointer from `mi_malloc`/the Rust global allocator must never
69///   be passed to [`unwrapped_free`]/[`unwrapped_realloc`]). Mixing these
70///   families corrupts allocator-internal bookkeeping.
71/// - The memory is uninitialized; reading it before writing is undefined
72///   behavior, as with any raw allocation.
73pub unsafe fn unwrapped_malloc(size: usize, alignment: usize) -> *mut u8 {
74    unsafe { sys::mi_unwrapped_malloc(size, alignment).cast() }
75}
76
77/// Free a pointer returned by [`unwrapped_malloc`] or [`unwrapped_realloc`].
78///
79/// Thin wrapper around `mi_unwrapped_free` (include/mimalloc/memory-events.h).
80///
81/// # Safety
82///
83/// - `p` must be either a null pointer (a documented, safe no-op on the C
84///   side) or a pointer previously returned by [`unwrapped_malloc`] or
85///   [`unwrapped_realloc`] that has not already been freed.
86/// - `p` must never have come from `mi_malloc`, this crate's [`MiMalloc`]
87///   allocator, or Rust's global allocator — passing such a pointer here is
88///   undefined behavior (the "unwrapped" and normal allocation families use
89///   incompatible header layouts and are validated by a magic-number check
90///   that a foreign pointer will not satisfy).
91pub unsafe fn unwrapped_free(p: *mut u8) {
92    unsafe { sys::mi_unwrapped_free(p.cast()) }
93}
94
95/// Resize a pointer returned by [`unwrapped_malloc`] or [`unwrapped_realloc`].
96///
97/// Thin wrapper around `mi_unwrapped_realloc` (include/mimalloc/memory-events.h).
98/// If `p` is null, this behaves like [`unwrapped_malloc`]. If `new_size` is
99/// `0`, this frees `p` (like [`unwrapped_free`]) and returns a null pointer.
100/// Otherwise the existing contents are copied into a freshly allocated
101/// unwrapped block (up to `min(old payload size, new_size)` bytes) and `p` is
102/// freed; `p` must not be used again after this call, whether or not it
103/// returns null.
104///
105/// Returns a null pointer on failure (including invalid `alignment`; see
106/// [`unwrapped_malloc`]'s `# Safety` section), in which case `p` is left
107/// valid and unfreed.
108///
109/// # Safety
110///
111/// - `p` must be either a null pointer or a pointer previously returned by
112///   [`unwrapped_malloc`] or [`unwrapped_realloc`] that has not already been
113///   freed, per the same family-isolation rule as [`unwrapped_free`].
114/// - `alignment` has the same power-of-two-or-zero contract as
115///   [`unwrapped_malloc`].
116/// - After this call, `p` must not be read, written, or freed again — treat
117///   it as consumed regardless of whether the return value is null.
118pub unsafe fn unwrapped_realloc(p: *mut u8, new_size: usize, alignment: usize) -> *mut u8 {
119    unsafe { sys::mi_unwrapped_realloc(p.cast(), new_size, alignment).cast() }
120}
121
122/// Turn on sampled heap profiling at the default sample rate.
123///
124/// Convenience entry point for wiring profiling to a command-line flag:
125///
126/// ```no_run
127/// # let args_profile_heap = true;
128/// if args_profile_heap {
129///     mimalloc_pprof::enable_heap_profiling();
130/// }
131/// ```
132///
133/// Uses the built-in default rate (one sample per ~512 KiB allocated;
134/// `MIMALLOC_PROF_SAMPLE_RATE` still overrides it). Call [`prof::start`]
135/// instead to pick a rate programmatically. Allocations made before this
136/// call — including process-startup and static initialization — are not
137/// tracked; profiles reflect steady-state behavior from this point on,
138/// which is the usual intent for an opt-in CLI switch. To capture startup
139/// as well, set `MIMALLOC_PROF=1` in the environment instead.
140///
141/// Returns `false` if profiling was already enabled (the earlier session,
142/// and its sample rate, stay active).
143pub fn enable_heap_profiling() -> bool {
144    prof::start(0)
145}
146
147/// How [`ProfConfig`] fields interact with the profiler's environment
148/// variables and `mi_option_*` settings.
149///
150/// Mirrors `mi_prof_config_mode_t` (include/mimalloc/profile.h); see that
151/// header for the full FALLBACK/OVERRIDE semantics, including the caveat
152/// that in `Override` mode `accum == false`, `dump_format == Text`, and
153/// `max_profiler_bytes == None` cannot be distinguished from "unset" and so
154/// always fall back to env-then-default rather than forcing the off/default
155/// value.
156#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
157pub enum ProfConfigMode {
158    /// Struct fields are used only where the corresponding env var / option is absent.
159    #[default]
160    Fallback,
161    /// Non-default struct fields win over env vars / options (see the caveat above).
162    Override,
163}
164
165/// Output format for [`ProfConfig::dump_at_exit`].
166///
167/// Mirrors `MI_PROF_FORMAT_TEXT` / `MI_PROF_FORMAT_PROTO` (include/mimalloc/profile.h).
168#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
169pub enum DumpFormat {
170    /// Legacy "heap profile:" text format (see [`prof::dump_to_vec`]).
171    #[default]
172    Text,
173    /// Binary pprof `profile.proto` format (see [`prof::dump_proto_to_vec`]).
174    Proto,
175}
176
177/// Ergonomic, Rust-facing sibling of `mi_prof_config_t`
178/// (include/mimalloc/profile.h) for [`enable_heap_profiling_with`].
179///
180/// Fields mirror the C struct one-for-one, but trade its 0/NULL-means-unset
181/// raw-integer conventions for `Option<T>` and enums where that reads
182/// better. `#[non_exhaustive]` + `Default` keeps future fields additive:
183/// build from `Default::default()` and set the fields you need, e.g.
184///
185/// ```
186/// use mimalloc_pprof::ProfConfig;
187/// let mut config = ProfConfig::default();
188/// config.sample_interval = Some(4096);
189/// ```
190///
191/// (Within this crate, struct-update syntax like
192/// `ProfConfig { sample_interval: Some(4096), ..Default::default() }` also
193/// works; `#[non_exhaustive]` only blocks struct-literal construction from
194/// *other* crates, so new fields stay non-breaking for them.)
195#[non_exhaustive]
196#[derive(Debug, Clone, Default)]
197pub struct ProfConfig {
198    /// See [`ProfConfigMode`].
199    pub mode: ProfConfigMode,
200    /// Average bytes between samples. `None` = env/default (512 KiB).
201    pub sample_interval: Option<usize>,
202    /// Budget (bytes) for profiler-internal persistent sampling state
203    /// (sample records, the stack intern table, interned stack entries).
204    /// `None` = unbudgeted (cap-bounded only).
205    pub max_profiler_bytes: Option<usize>,
206    /// `None` = nondeterministic.
207    pub seed: Option<u64>,
208    pub accum: bool,
209    /// `None` = default (32); compile cap 128.
210    pub max_stack_depth: Option<usize>,
211    /// Path to dump the profile to at process exit. `None` = no exit dump.
212    pub dump_at_exit: Option<PathBuf>,
213    /// Format used for the exit dump. Ignored if `dump_at_exit` is `None`.
214    pub dump_format: DumpFormat,
215}
216
217/// Turn on sampled heap profiling using a struct-based configuration.
218///
219/// Sibling of [`enable_heap_profiling`] for callers that need more than a
220/// single sample rate -- e.g. seeding the sampler, capping profiler-arena
221/// memory, or registering an exit-time dump path/format. See [`ProfConfig`]
222/// and, for the full FALLBACK/OVERRIDE semantics, `mi_prof_config_mode_t` in
223/// `include/mimalloc/profile.h`.
224///
225/// Returns `false` if profiling was already enabled (the earlier session
226/// stays active), or if `config.dump_at_exit` is set but is not
227/// representable as a NUL-free C string (non-UTF-8 or an embedded NUL byte)
228/// -- in that case `mi_prof_start_ex` is never called.
229pub fn enable_heap_profiling_with(config: &ProfConfig) -> bool {
230    // `dump_at_exit_c` must outlive the `mi_prof_start_ex` call below since
231    // `raw.dump_at_exit` borrows its bytes; it does, as both live to the end
232    // of this function.
233    let dump_at_exit_c: Option<CString> = match &config.dump_at_exit {
234        Some(path) => match path.to_str().and_then(|s| CString::new(s).ok()) {
235            Some(c) => Some(c),
236            None => return false,
237        },
238        None => None,
239    };
240
241    let mut raw: sys::mi_prof_config_t = unsafe { core::mem::zeroed() };
242    raw.size = core::mem::size_of::<sys::mi_prof_config_t>();
243    raw.version = sys::MI_PROF_CONFIG_VERSION;
244    raw.mode = match config.mode {
245        ProfConfigMode::Fallback => sys::MI_PROF_CONFIG_FALLBACK,
246        ProfConfigMode::Override => sys::MI_PROF_CONFIG_OVERRIDE,
247    };
248    raw.sample_interval = config.sample_interval.unwrap_or(0);
249    raw.max_profiler_bytes = config.max_profiler_bytes.unwrap_or(0);
250    raw.seed = config.seed.unwrap_or(0);
251    raw.accum = config.accum;
252    raw.max_stack_depth = config.max_stack_depth.unwrap_or(0);
253    raw.dump_at_exit = dump_at_exit_c
254        .as_ref()
255        .map_or(core::ptr::null(), |c| c.as_ptr());
256    raw.dump_format = match config.dump_format {
257        DumpFormat::Text => sys::MI_PROF_FORMAT_TEXT,
258        DumpFormat::Proto => sys::MI_PROF_FORMAT_PROTO,
259    };
260
261    unsafe { sys::mi_prof_start_ex(&raw) }
262}
263
264/// Safe controls for mimalloc's sampled heap profiler.
265pub mod prof {
266    use core::ffi::{c_char, c_void};
267    use std::ffi::{CStr, CString};
268    use std::io;
269    use std::panic::{catch_unwind, AssertUnwindSafe};
270    use std::path::Path;
271
272    use crate::sys;
273
274    pub fn start(sample_rate: usize) -> bool {
275        unsafe { sys::mi_prof_start(sample_rate) }
276    }
277    #[doc(hidden)]
278    pub fn start_seeded(sample_rate: usize, seed: u64) -> bool {
279        unsafe { sys::mi_prof_start_seeded(sample_rate, seed) }
280    }
281    pub fn stop() {
282        unsafe { sys::mi_prof_stop() }
283    }
284    pub fn is_enabled() -> bool {
285        unsafe { sys::mi_prof_is_enabled() }
286    }
287    pub fn reset() {
288        unsafe { sys::mi_prof_reset() }
289    }
290
291    pub fn dump_file(path: &Path) -> io::Result<()> {
292        let path = path.to_str().ok_or_else(|| {
293            io::Error::new(io::ErrorKind::InvalidInput, "profile path is not UTF-8")
294        })?;
295        let path = CString::new(path).map_err(|_| {
296            io::Error::new(io::ErrorKind::InvalidInput, "profile path contains NUL")
297        })?;
298        if unsafe { sys::mi_prof_dump(path.as_ptr()) } {
299            Ok(())
300        } else {
301            Err(io::Error::last_os_error())
302        }
303    }
304
305    unsafe extern "C" fn write_cb(arg: *mut c_void, buf: *const c_char, len: usize) {
306        let out = &mut *(arg as *mut Vec<u8>);
307        out.extend_from_slice(core::slice::from_raw_parts(buf.cast::<u8>(), len));
308    }
309
310    /// Serialize the current heap profile without holding the profiler lock.
311    pub fn dump_to_vec() -> Vec<u8> {
312        let mut out = Vec::new();
313        let ok =
314            unsafe { sys::mi_prof_dump_writer(Some(write_cb), (&mut out as *mut Vec<u8>).cast()) };
315        if ok {
316            out
317        } else {
318            Vec::new()
319        }
320    }
321
322    /// Serialize the current heap profile as a binary pprof `profile.proto`
323    /// `Profile` message (see [google/pprof's `profile.proto`][proto]),
324    /// without holding the profiler lock.
325    ///
326    /// Sample values are pre-scaled the same way Go's `runtime/pprof` scales
327    /// legacy heap samples (the `protomem.go` convention: `alloc_objects`,
328    /// `alloc_space`, `inuse_objects`, `inuse_space`, each already corrected
329    /// for Poisson sampling bias rather than left for a downstream tool to
330    /// rescale). The `Mapping` table is included, so external symbolizers
331    /// need only the binary — no text parsing of a "heap profile:" header or
332    /// a `MAPPED_LIBRARIES:` section. This is the compact, machine-oriented
333    /// counterpart to [`dump_to_vec`]'s text format, intended for API and
334    /// transport use (issue #23) where a `pprof`-compatible tool consumes
335    /// the bytes directly.
336    ///
337    /// [proto]: https://github.com/google/pprof/blob/main/proto/profile.proto
338    pub fn dump_proto_to_vec() -> Vec<u8> {
339        let mut out = Vec::new();
340        let ok = unsafe {
341            sys::mi_prof_dump_proto_writer(Some(write_cb), (&mut out as *mut Vec<u8>).cast())
342        };
343        if ok {
344            out
345        } else {
346            Vec::new()
347        }
348    }
349
350    /// Write the current heap profile to `path` in `profile.proto` format.
351    ///
352    /// See [`dump_proto_to_vec`] for the format details.
353    pub fn dump_proto_file(path: &Path) -> io::Result<()> {
354        let path = path.to_str().ok_or_else(|| {
355            io::Error::new(io::ErrorKind::InvalidInput, "profile path is not UTF-8")
356        })?;
357        let path = CString::new(path).map_err(|_| {
358            io::Error::new(io::ErrorKind::InvalidInput, "profile path contains NUL")
359        })?;
360        if unsafe { sys::mi_prof_dump_proto(path.as_ptr()) } {
361            Ok(())
362        } else {
363            Err(io::Error::last_os_error())
364        }
365    }
366
367    /// Snapshot of `mi_prof_stats_get`'s counters, translated from the raw
368    /// sys struct into plain Rust types.
369    #[derive(Debug, Clone, Default)]
370    pub struct ProfStats {
371        pub enabled: bool,
372        pub accum: bool,
373        pub sample_rate: usize,
374        pub live_samples: usize,
375        pub live_bytes: usize,
376        pub accum_samples: usize,
377        pub accum_bytes: usize,
378        pub unique_stacks: usize,
379        pub arena_committed: usize,
380        pub stack_table_overflows: usize,
381        /// Count of ALL dropped samples (record-alloc failure, stack-intern
382        /// failure, including the stack-table cap); a superset of
383        /// `stack_table_overflows`, so `dropped_samples >=
384        /// stack_table_overflows` always.
385        pub dropped_samples: usize,
386        /// Allocator-level ("ground truth") counters, read from the mimalloc v3
387        /// engine's per-heap statistics at the time of the call. Every field
388        /// above is *sampled*; these are exact, so comparing them against
389        /// `live_bytes` measures the sampler's error directly -- which is what
390        /// makes an assertion on a sampled profile meaningful in a test.
391        pub heap: HeapStats,
392    }
393
394    /// Exact allocator counters accompanying a [`ProfStats`] reading.
395    ///
396    /// These come from mimalloc v3's per-heap statistics
397    /// (`mi_heap_stats_get`/`mi_subproc_stats_get`), which the v2 engine did not
398    /// expose. They are valid even when the profiler is stopped.
399    #[derive(Debug, Clone, Default)]
400    pub struct HeapStats {
401        /// Bytes currently committed from the OS.
402        pub committed: usize,
403        /// Bytes currently reserved from the OS (always `>= committed`).
404        pub reserved: usize,
405        /// Bytes the application actually requested and still holds.
406        ///
407        /// Only maintained when the C library was built with `MI_STAT >= 2`;
408        /// otherwise this is 0. Check [`HeapStats::detailed`] before using it.
409        pub malloc_requested: usize,
410        /// Live mimalloc pages.
411        pub pages: usize,
412        /// Pages abandoned by exited threads.
413        pub pages_abandoned: usize,
414        /// Live first-class heaps.
415        pub heaps: usize,
416        /// Live thread-local heaps. The main thread's statically-initialized
417        /// theap is not counted, so a single-threaded process reports 0.
418        pub theaps: usize,
419        /// Cumulative bytes purged back to the OS.
420        pub purged: usize,
421        /// Whether the C library was built with `MI_STAT >= 2` ("detailed"
422        /// statistics), which upstream enables by default only for debug
423        /// builds. [`HeapStats::malloc_requested`] is maintained only at that
424        /// level; every other field here is maintained at any level.
425        ///
426        /// Without this flag you cannot tell "the application allocated
427        /// nothing" from "this build does not track that counter".
428        pub detailed: bool,
429    }
430
431    /// Read the profiler's current counters via `mi_prof_stats_get`.
432    ///
433    /// Returns `ProfStats::default()` (all zero/false) if the call fails,
434    /// e.g. because the sys struct's `size`/`version` header does not match
435    /// what the linked mimalloc build expects.
436    pub fn stats() -> ProfStats {
437        let mut raw: sys::mi_prof_stats_t = unsafe { core::mem::zeroed() };
438        raw.size = core::mem::size_of::<sys::mi_prof_stats_t>();
439        raw.version = sys::MI_PROF_STAT_VERSION;
440        if unsafe { sys::mi_prof_stats_get(&mut raw) } {
441            ProfStats {
442                enabled: raw.enabled,
443                accum: raw.accum,
444                sample_rate: raw.sample_rate,
445                live_samples: raw.live_samples,
446                live_bytes: raw.live_bytes,
447                accum_samples: raw.accum_samples,
448                accum_bytes: raw.accum_bytes,
449                unique_stacks: raw.unique_stacks,
450                arena_committed: raw.arena_committed,
451                stack_table_overflows: raw.stack_table_overflows,
452                dropped_samples: raw.dropped_samples,
453                heap: HeapStats {
454                    committed: raw.heap_committed,
455                    reserved: raw.heap_reserved,
456                    malloc_requested: raw.heap_malloc_requested,
457                    pages: raw.heap_pages,
458                    pages_abandoned: raw.heap_pages_abandoned,
459                    heaps: raw.heap_count,
460                    theaps: raw.theap_count,
461                    purged: raw.heap_purged,
462                    detailed: raw.heap_stats_detailed,
463                },
464            }
465        } else {
466            ProfStats::default()
467        }
468    }
469
470    /// One sampled call stack, copied out of the profiler by [`samples`].
471    #[derive(Debug, Clone)]
472    pub struct Sample {
473        pub stack: Vec<usize>,
474        pub live_objects: usize,
475        pub live_bytes: usize,
476        pub accum_objects: usize,
477        pub accum_bytes: usize,
478    }
479
480    impl Sample {
481        /// Estimate the un-sampled byte volume behind this sample.
482        ///
483        /// Mirrors pprof's legacy heap-sample scaling formula
484        /// (`scaleHeapSample` in pprof's `profile/legacy_profile.go`),
485        /// which corrects for the bias a Poisson sampling process with mean
486        /// interval `sample_rate` introduces toward larger allocations.
487        pub fn estimated_bytes(&self, sample_rate: usize) -> u64 {
488            if self.live_objects == 0 || self.live_bytes == 0 {
489                return 0;
490            }
491            if sample_rate <= 1 {
492                return self.live_bytes as u64;
493            }
494            let avg = self.live_bytes as f64 / self.live_objects as f64;
495            let scale = 1.0 / (1.0 - (-avg / sample_rate as f64).exp());
496            (self.live_bytes as f64 * scale) as u64
497        }
498    }
499
500    /// Frees the snapshot handle on drop, including on unwind, so a panic
501    /// partway through collection never leaks profiler-arena memory.
502    struct SnapshotGuard(*mut sys::mi_prof_snapshot_t);
503
504    impl Drop for SnapshotGuard {
505        fn drop(&mut self) {
506            unsafe { sys::mi_prof_snapshot_free(self.0) }
507        }
508    }
509
510    unsafe extern "C" fn collect_visitor(
511        info: *const sys::mi_prof_sample_info_t,
512        arg: *mut c_void,
513    ) -> bool {
514        let result = catch_unwind(AssertUnwindSafe(|| unsafe {
515            let out = &mut *(arg as *mut Vec<Sample>);
516            let info = &*info;
517            let stack = (0..info.depth)
518                .map(|i| *info.stack.add(i) as usize)
519                .collect();
520            out.push(Sample {
521                stack,
522                live_objects: info.live_objects,
523                live_bytes: info.live_bytes,
524                accum_objects: info.accum_objects,
525                accum_bytes: info.accum_bytes,
526            });
527        }));
528        result.is_ok()
529    }
530
531    /// Collect a point-in-time copy of every live sampled stack.
532    ///
533    /// This snapshots under the profiler lock via `mi_prof_snapshot_new`,
534    /// then walks and frees the snapshot outside that lock. Using
535    /// `mi_prof_visit` directly here would run the (allocating) collection
536    /// below from inside the visitor while the profiler lock is held,
537    /// risking reentrant profiler-hook allocation and deadlock — the
538    /// reentrancy hazard the snapshot API exists to avoid (issue #2,
539    /// decisions 11-13).
540    pub fn samples() -> Vec<Sample> {
541        let snap = unsafe { sys::mi_prof_snapshot_new() };
542        if snap.is_null() {
543            return Vec::new();
544        }
545        let guard = SnapshotGuard(snap);
546        let mut out: Vec<Sample> = Vec::new();
547        unsafe {
548            sys::mi_prof_snapshot_visit(
549                guard.0,
550                collect_visitor,
551                (&mut out as *mut Vec<Sample>).cast(),
552            );
553        }
554        out
555    }
556
557    /// One loaded module (shared library or the main executable), copied out
558    /// of the OS module list by [`modules`].
559    #[derive(Debug, Clone)]
560    pub struct ModuleInfo {
561        pub path: String,
562        pub base: usize,
563        pub size: usize,
564    }
565
566    unsafe extern "C" fn modules_visitor(
567        info: *const sys::mi_prof_module_info_t,
568        arg: *mut c_void,
569    ) -> bool {
570        let result = catch_unwind(AssertUnwindSafe(|| unsafe {
571            let out = &mut *(arg as *mut Vec<ModuleInfo>);
572            let info = &*info;
573            // `info.path` is only valid for the duration of this callback (it
574            // points into OS-owned module-list storage), so it must be copied
575            // into an owned `String` right here rather than stashed for later.
576            let path = CStr::from_ptr(info.path).to_string_lossy().into_owned();
577            out.push(ModuleInfo {
578                path,
579                base: info.base,
580                size: info.size,
581            });
582        }));
583        result.is_ok()
584    }
585
586    /// Enumerate the process's loaded modules (shared libraries and the main
587    /// executable), e.g. to build pprof `Mapping` entries yourself.
588    ///
589    /// Unlike [`samples`]'s `collect_visitor`, this callback is free to
590    /// allocate: `mi_prof_modules_visit` never takes the profiler lock (the
591    /// module list is OS-owned, not part of the sampled-allocation table), so
592    /// there is no reentrant-allocation-under-the-lock hazard here.
593    pub fn modules() -> Vec<ModuleInfo> {
594        let mut out: Vec<ModuleInfo> = Vec::new();
595        unsafe {
596            sys::mi_prof_modules_visit(
597                modules_visitor,
598                (&mut out as *mut Vec<ModuleInfo>).cast(),
599            );
600        }
601        out
602    }
603}
604
605#[cfg(test)]
606mod tests {
607    use super::*;
608    use std::sync::Mutex;
609
610    // The profiler is process-global state, and unit tests within this
611    // binary may run concurrently by default, so serialize everything that
612    // starts/stops it. `unwrap_or_else` rides through a poisoned lock rather
613    // than cascading a single panicking test into every other one.
614    static PROF_TEST_LOCK: Mutex<()> = Mutex::new(());
615
616    fn reset_profiler() {
617        if prof::is_enabled() {
618            prof::stop();
619        }
620    }
621
622    #[test]
623    fn enable_heap_profiling_with_default_config_starts_profiler() {
624        let _guard = PROF_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
625        reset_profiler();
626
627        let config = ProfConfig::default();
628        assert!(enable_heap_profiling_with(&config));
629        assert!(prof::is_enabled());
630
631        prof::stop();
632    }
633
634    #[test]
635    fn enable_heap_profiling_with_override_mode_sets_sample_interval() {
636        let _guard = PROF_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
637        reset_profiler();
638
639        let config = ProfConfig {
640            mode: ProfConfigMode::Override,
641            sample_interval: Some(4096),
642            ..Default::default()
643        };
644        assert!(enable_heap_profiling_with(&config));
645        assert!(prof::is_enabled());
646        assert_eq!(prof::stats().sample_rate, 4096);
647
648        prof::stop();
649    }
650
651    #[test]
652    fn unwrapped_malloc_write_realloc_grow_verify_free() {
653        unsafe {
654            let size = 64usize;
655            let p = unwrapped_malloc(size, 0);
656            assert!(!p.is_null());
657
658            for i in 0..size {
659                *p.add(i) = (i % 256) as u8;
660            }
661
662            let new_size = 256usize;
663            let p2 = unwrapped_realloc(p, new_size, 0);
664            assert!(!p2.is_null());
665
666            for i in 0..size {
667                assert_eq!(*p2.add(i), (i % 256) as u8);
668            }
669
670            unwrapped_free(p2);
671        }
672    }
673
674    #[test]
675    fn unwrapped_free_null_is_noop() {
676        unsafe {
677            unwrapped_free(core::ptr::null_mut());
678        }
679    }
680
681    #[test]
682    fn unwrapped_malloc_rejects_non_power_of_two_alignment() {
683        unsafe {
684            let p = unwrapped_malloc(16, 3);
685            assert!(p.is_null());
686        }
687    }
688
689    #[test]
690    fn unwrapped_realloc_with_null_ptr_behaves_like_malloc() {
691        unsafe {
692            let p = unwrapped_realloc(core::ptr::null_mut(), 32, 0);
693            assert!(!p.is_null());
694            unwrapped_free(p);
695        }
696    }
697
698    #[test]
699    fn unwrapped_realloc_with_zero_size_frees_and_returns_null() {
700        unsafe {
701            let p = unwrapped_malloc(32, 0);
702            assert!(!p.is_null());
703            let p2 = unwrapped_realloc(p, 0, 0);
704            assert!(p2.is_null());
705        }
706    }
707}