Skip to main content

pprof_alloc/
lib.rs

1//! Allocation profiling and Linux memory telemetry for Rust services.
2//!
3//! `pprof-alloc` provides a [`GlobalAlloc`] wrapper that can sample allocation
4//! stack traces and export them as gzipped pprof heap profiles. It also exposes
5//! Linux memory collectors for allocator state, cgroup v2 memory accounting, and
6//! `/proc/self/smaps_rollup` process residency.
7//!
8//! The crate is intended to be embedded in binaries that expose their own debug
9//! or metrics endpoint. Use [`PprofAlloc`] as the process global allocator, then
10//! call [`generate_pprof`] or [`snapshot`] from your application surface.
11//!
12//! [`GlobalAlloc`]: std::alloc::GlobalAlloc
13
14pub mod allocator;
15mod env;
16mod pprof;
17pub mod stats;
18mod trace;
19
20pub use crate::env::{ALLOCATOR_ENV, Allocator, PPROF_BACKEND_ENV, PPROF_SAMPLE_RATE_ENV};
21use crate::env::{AllocatorSelection, PprofBackend};
22use crate::pprof::{StackProfile, WeightedStack};
23use crate::trace::HashedBacktrace;
24use dashmap::DashMap;
25use serde::Serialize;
26use std::alloc::{GlobalAlloc, Layout, System};
27use std::cell::Cell;
28use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
29use std::time::{SystemTime, UNIX_EPOCH};
30
31pub use crate::trace::CaptureMode;
32
33/// Default average number of allocated bytes between recorded pprof samples.
34///
35/// This matches Go's default heap profiling rate: one sampled allocation per
36/// 512 KiB of allocated bytes on average. A rate of `1` records every
37/// allocation, while `0` disables pprof stack recording.
38pub const DEFAULT_PPROF_SAMPLE_RATE: usize = 512 * 1024;
39
40const MAX_FAST_EXP_RAND_MEAN: usize = 0x7000000;
41const RANDOM_BIT_COUNT: u32 = 26;
42const RESOLVED_SAMPLE_RATE_UNINITIALIZED: usize = usize::MAX;
43const STATS_FLUSH_EVENTS: u64 = 1024;
44const STATS_FLUSH_BYTES: u64 = 1024 * 1024;
45
46#[repr(u8)]
47#[derive(Clone, Copy, Debug, PartialEq, Eq)]
48enum TrackingMode {
49	Uninitialized = 0,
50	System = 1,
51	Jemalloc = 2,
52	Mimalloc = 3,
53	Stats = 4,
54	Pprof = 5,
55	PprofStats = 6,
56}
57
58impl TrackingMode {
59	const fn from_u8(value: u8) -> Self {
60		match value {
61			1 => Self::System,
62			2 => Self::Jemalloc,
63			3 => Self::Mimalloc,
64			4 => Self::Stats,
65			5 => Self::Pprof,
66			6 => Self::PprofStats,
67			_ => Self::Uninitialized,
68		}
69	}
70}
71
72/// Global allocator that can collect allocation counters and pprof heap profiles.
73///
74/// Use this as the process `#[global_allocator]`. The backing allocator is
75/// selected from [`ALLOCATOR_ENV`] on first allocation.
76pub struct PprofAlloc {
77	/// Enable profiling support
78	pprof: bool,
79	/// Enable coarse grained stats
80	stats: bool,
81	/// Average bytes between pprof samples. 0 disables pprof, 1 records everything.
82	pprof_sample_rate: usize,
83	/// Read the pprof sample rate from PPROF_ALLOC_SAMPLE_RATE at runtime.
84	pprof_sample_rate_from_env: bool,
85	/// Cached resolved sample rate for env-driven configuration.
86	resolved_pprof_sample_rate: AtomicUsize,
87	/// Cached wrapper work needed on allocation/deallocation.
88	tracking_mode: AtomicU8,
89	/// Allocator to use when [`ALLOCATOR_ENV`] is unset.
90	default_allocator: Allocator,
91}
92
93#[derive(Clone)]
94struct AllocationRecord {
95	size: usize,
96	trace: HashedBacktrace,
97}
98
99struct HeapSampleValues {
100	alloc_objects: i64,
101	alloc_space: i64,
102	inuse_objects: i64,
103	inuse_space: i64,
104}
105
106struct LocalAllocationStats {
107	allocated: Cell<u64>,
108	freed: Cell<u64>,
109	allocations: Cell<u64>,
110	frees: Cell<u64>,
111}
112
113impl LocalAllocationStats {
114	const fn new() -> Self {
115		Self {
116			allocated: Cell::new(0),
117			freed: Cell::new(0),
118			allocations: Cell::new(0),
119			frees: Cell::new(0),
120		}
121	}
122
123	fn record_allocation(&self, size: u64) {
124		self
125			.allocated
126			.set(self.allocated.get().saturating_add(size));
127		self
128			.allocations
129			.set(self.allocations.get().saturating_add(1));
130		self.flush_if_needed();
131	}
132
133	fn record_deallocation(&self, size: u64) {
134		self.freed.set(self.freed.get().saturating_add(size));
135		self.frees.set(self.frees.get().saturating_add(1));
136		self.flush_if_needed();
137	}
138
139	fn flush_if_needed(&self) {
140		let events = self.allocations.get().saturating_add(self.frees.get());
141		let bytes = self.allocated.get().saturating_add(self.freed.get());
142		if events >= STATS_FLUSH_EVENTS || bytes >= STATS_FLUSH_BYTES {
143			self.flush();
144		}
145	}
146
147	fn flush(&self) {
148		let allocated = self.allocated.replace(0);
149		let freed = self.freed.replace(0);
150		let allocations = self.allocations.replace(0);
151		let frees = self.frees.replace(0);
152
153		if allocated != 0 {
154			GLOBAL_STATS
155				.allocated
156				.fetch_add(allocated, Ordering::Relaxed);
157		}
158		if freed != 0 {
159			GLOBAL_STATS.freed.fetch_add(freed, Ordering::Relaxed);
160		}
161		if allocations != 0 {
162			GLOBAL_STATS
163				.allocations
164				.fetch_add(allocations, Ordering::Relaxed);
165		}
166		if frees != 0 {
167			GLOBAL_STATS.frees.fetch_add(frees, Ordering::Relaxed);
168		}
169	}
170
171	#[cfg(test)]
172	fn reset(&self) {
173		self.allocated.set(0);
174		self.freed.set(0);
175		self.allocations.set(0);
176		self.frees.set(0);
177	}
178}
179
180impl Drop for LocalAllocationStats {
181	fn drop(&mut self) {
182		self.flush();
183	}
184}
185
186impl HeapSampleValues {
187	fn from_allocations(stats: &stats::Allocations, sample_rate: usize) -> Self {
188		let (alloc_objects, alloc_space) =
189			scale_heap_sample(stats.allocations, stats.allocated, sample_rate);
190		let (inuse_objects, inuse_space) = scale_heap_sample(
191			stats.in_use_allocations(),
192			stats.in_use_bytes(),
193			sample_rate,
194		);
195		Self {
196			alloc_objects,
197			alloc_space,
198			inuse_objects,
199			inuse_space,
200		}
201	}
202}
203
204fn saturating_i64(value: u64) -> i64 {
205	value.min(i64::MAX as u64) as i64
206}
207
208fn scale_heap_sample(count: u64, size: u64, sample_rate: usize) -> (i64, i64) {
209	if count == 0 || size == 0 {
210		return (0, 0);
211	}
212
213	if sample_rate <= 1 {
214		return (saturating_i64(count), saturating_i64(size));
215	}
216
217	let average_size = size as f64 / count as f64;
218	let probability = -(-average_size / sample_rate as f64).exp_m1();
219	if probability <= 0.0 {
220		return (saturating_i64(count), saturating_i64(size));
221	}
222	let scale = 1.0 / probability;
223	(
224		(count as f64 * scale).min(i64::MAX as f64) as i64,
225		(size as f64 * scale).min(i64::MAX as f64) as i64,
226	)
227}
228
229#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize)]
230/// Summary of the currently recorded pprof allocation profile.
231///
232/// Values are scaled according to the active pprof sample rate, so they are
233/// estimates when sampling is enabled.
234pub struct PprofSummary {
235	/// Number of distinct stack traces recorded in the profile.
236	pub total_stacks: u64,
237	/// Number of distinct stack traces with estimated live bytes greater than zero.
238	pub live_stacks: u64,
239	/// Estimated cumulative allocated bytes across all recorded stacks.
240	pub alloc_space_bytes: u64,
241	/// Estimated currently live bytes across all recorded stacks.
242	pub inuse_space_bytes: u64,
243	/// Estimated cumulative allocation count across all recorded stacks.
244	pub alloc_objects: u64,
245	/// Estimated currently live allocation count across all recorded stacks.
246	pub inuse_objects: u64,
247}
248
249#[derive(Clone, Debug, Serialize)]
250/// Best-effort snapshot of allocation and memory state for the current process created by `snapshot()`.
251///
252/// Snapshot collection never fails as a whole. Optional fields are `None` when
253/// the corresponding operating-system or allocator probe is unavailable.
254pub struct MemorySnapshot {
255	/// Wall-clock capture time as milliseconds since the Unix epoch.
256	pub captured_at_unix_ms: u64,
257	/// Stack capture implementation compiled into this build.
258	pub capture_mode: CaptureMode,
259	/// Coarse process-wide counters from [`allocation_stats`].
260	pub allocation_stats: stats::Allocations,
261	/// Summary of recorded pprof stack-attributed allocation data.
262	pub pprof: PprofSummary,
263	/// Active allocator identity and allocator-specific memory stats, if available.
264	pub allocator: allocator::AllocatorSnapshot,
265	/// cgroup v2 memory stats for this process, when available.
266	pub cgroup: Option<stats::cgroups::MemoryStat>,
267	/// `/proc/self/smaps_rollup` memory stats for this process, when available.
268	pub smaps: Option<stats::smaps::ProcessStats>,
269}
270
271impl Default for PprofAlloc {
272	fn default() -> Self {
273		Self::new()
274	}
275}
276
277impl PprofAlloc {
278	/// Create an allocator with profiling disabled.
279	///
280	/// Use [`Self::with_pprof`], [`Self::with_pprof_sample_rate`], or
281	/// [`Self::with_stats`] to enable collection.
282	pub const fn new() -> Self {
283		PprofAlloc {
284			pprof: false,
285			stats: false,
286			pprof_sample_rate: DEFAULT_PPROF_SAMPLE_RATE,
287			pprof_sample_rate_from_env: false,
288			resolved_pprof_sample_rate: AtomicUsize::new(RESOLVED_SAMPLE_RATE_UNINITIALIZED),
289			tracking_mode: AtomicU8::new(TrackingMode::Uninitialized as u8),
290			default_allocator: Allocator::System,
291		}
292	}
293
294	/// Set the allocator used when [`ALLOCATOR_ENV`] is unset.
295	pub const fn with_default(mut self, allocator: Allocator) -> Self {
296		self.default_allocator = allocator;
297		self
298	}
299
300	/// Enable sampled pprof stack profiling using [`DEFAULT_PPROF_SAMPLE_RATE`].
301	///
302	/// When the active allocator is jemalloc and the backend env selects native
303	/// profiling, builds with `allocator-jemalloc` use jemalloc's native heap
304	/// profiler instead of wrapper-side allocation tracking.
305	pub const fn with_pprof(mut self) -> Self {
306		self.pprof = true;
307		self
308	}
309
310	/// Enable sampled pprof stack profiling with an explicit byte sample rate.
311	///
312	/// A rate of `1` records every allocation. A rate of `0` disables pprof
313	/// stack recording while still allowing other enabled collectors to run.
314	/// When native jemalloc profiling is selected, this wrapper-side rate is
315	/// ignored. Use [`PPROF_SAMPLE_RATE_ENV`] with [`configure`] to set jemalloc's
316	/// runtime sample rate.
317	pub const fn with_pprof_sample_rate(mut self, bytes: usize) -> Self {
318		self.pprof = true;
319		self.pprof_sample_rate = bytes;
320		self.pprof_sample_rate_from_env = false;
321		self
322	}
323
324	/// Enable pprof stack profiling with the sample rate read from the environment.
325	///
326	/// [`PPROF_SAMPLE_RATE_ENV`] is read lazily on the first profiled allocation.
327	/// If the variable is missing or invalid, `default_rate` is used.
328	/// When native jemalloc profiling is selected, [`configure`] applies this
329	/// environment variable to jemalloc's runtime sample rate.
330	pub const fn with_pprof_sample_rate_from_env(mut self, default_rate: usize) -> Self {
331		self.pprof = true;
332		self.pprof_sample_rate = default_rate;
333		self.pprof_sample_rate_from_env = true;
334		self
335	}
336
337	/// Enable coarse process-wide allocation and free counters.
338	///
339	/// When the active allocator is jemalloc and jemalloc support is compiled
340	/// in, snapshots use jemalloc's native process stats instead of these
341	/// wrapper-side counters.
342	pub const fn with_stats(mut self) -> Self {
343		self.stats = true;
344		self
345	}
346
347	fn effective_pprof_sample_rate(&self) -> usize {
348		if self.pprof_sample_rate_from_env {
349			let resolved = self.resolved_pprof_sample_rate.load(Ordering::Relaxed);
350			if resolved != RESOLVED_SAMPLE_RATE_UNINITIALIZED {
351				return resolved;
352			}
353
354			let resolved = env_pprof_sample_rate(self.pprof_sample_rate);
355			self
356				.resolved_pprof_sample_rate
357				.store(resolved, Ordering::Relaxed);
358			resolved
359		} else {
360			self.pprof_sample_rate
361		}
362	}
363
364	#[cfg(test)]
365	fn active_pprof_sample_rate(&self) -> Option<usize> {
366		if !matches!(
367			self.tracking_mode(),
368			TrackingMode::Pprof | TrackingMode::PprofStats
369		) {
370			return None;
371		}
372
373		let sample_rate = self.effective_pprof_sample_rate();
374		CURRENT_PPROF_SAMPLE_RATE.store(sample_rate, Ordering::Relaxed);
375		(sample_rate != 0).then_some(sample_rate)
376	}
377
378	fn record_allocation_stats(&self, size: usize) {
379		if LOCAL_STATS
380			.try_with(|stats| stats.record_allocation(size as u64))
381			.is_err()
382		{
383			GLOBAL_STATS
384				.allocated
385				.fetch_add(size as u64, Ordering::Relaxed);
386			GLOBAL_STATS.allocations.fetch_add(1, Ordering::Relaxed);
387		}
388	}
389
390	fn record_deallocation_stats(&self, size: usize) {
391		if LOCAL_STATS
392			.try_with(|stats| stats.record_deallocation(size as u64))
393			.is_err()
394		{
395			GLOBAL_STATS.freed.fetch_add(size as u64, Ordering::Relaxed);
396			GLOBAL_STATS.frees.fetch_add(1, Ordering::Relaxed);
397		}
398	}
399
400	#[cfg(test)]
401	fn record_allocation(&self, ptr: usize, size: usize) {
402		if self.stats {
403			self.record_allocation_stats(size);
404		}
405		if let Some(sample_rate) = self.active_pprof_sample_rate() {
406			self.record_profile_allocation(ptr, size, sample_rate);
407		}
408	}
409
410	fn record_profile_allocation(&self, ptr: usize, size: usize, sample_rate: usize) {
411		if should_sample_allocation(size, sample_rate) {
412			enter_alloc(|| {
413				let trace = HashedBacktrace::capture();
414				self.record_allocation_with_trace(ptr, size, trace);
415			});
416		}
417	}
418
419	fn record_tracked_allocation(
420		&self,
421		ptr: *mut u8,
422		size: usize,
423		sample_rate: Option<usize>,
424		track_stats: bool,
425	) {
426		if ptr.is_null() {
427			return;
428		}
429
430		if track_stats {
431			self.record_allocation_stats(size);
432		}
433		if let Some(sample_rate) = sample_rate {
434			self.record_profile_allocation(ptr as usize, size, sample_rate);
435		}
436	}
437
438	fn tracking_is_recursive(&self, sample_rate: Option<usize>, track_stats: bool) -> bool {
439		(sample_rate.is_some() || track_stats) && IN_ALLOC.try_with(|x| x.get()).unwrap_or(true)
440	}
441
442	fn tracking_mode(&self) -> TrackingMode {
443		let mode = TrackingMode::from_u8(self.tracking_mode.load(Ordering::Relaxed));
444		if mode != TrackingMode::Uninitialized {
445			return mode;
446		}
447
448		let selected_allocator = env::selected_allocator(self.default_allocator);
449		let native_jemalloc =
450			cfg!(feature = "allocator-jemalloc") && selected_allocator == AllocatorSelection::Jemalloc;
451		let native_pprof = native_jemalloc && env::selected_pprof_backend() == PprofBackend::Native;
452		let wrapper_stats = self.stats && !native_jemalloc;
453		let wrapper_pprof_rate = (self.pprof && !native_pprof).then(|| {
454			let sample_rate = self.effective_pprof_sample_rate();
455			CURRENT_PPROF_SAMPLE_RATE.store(sample_rate, Ordering::Relaxed);
456			sample_rate
457		});
458		let wrapper_pprof = wrapper_pprof_rate.is_some_and(|sample_rate| sample_rate != 0);
459
460		let mode = match (wrapper_pprof, wrapper_stats) {
461			(false, false) => match selected_allocator {
462				AllocatorSelection::Jemalloc => TrackingMode::Jemalloc,
463				AllocatorSelection::Mimalloc => TrackingMode::Mimalloc,
464				_ => TrackingMode::System,
465			},
466			(false, true) => TrackingMode::Stats,
467			(true, false) => TrackingMode::Pprof,
468			(true, true) => TrackingMode::PprofStats,
469		};
470		self.tracking_mode.store(mode as u8, Ordering::Relaxed);
471		mode
472	}
473
474	#[cfg(all(test, feature = "allocator-jemalloc"))]
475	fn native_pprof_selected(&self) -> bool {
476		match env::selected_pprof_backend() {
477			PprofBackend::Wrapper => false,
478			PprofBackend::Native => {
479				env::selected_allocator(self.default_allocator) == AllocatorSelection::Jemalloc
480			},
481			PprofBackend::Uninitialized => false,
482		}
483	}
484
485	fn record_allocation_with_trace(&self, ptr: usize, size: usize, trace: HashedBacktrace) {
486		POINTER_MAP.insert(
487			ptr,
488			AllocationRecord {
489				size,
490				trace: trace.clone(),
491			},
492		);
493		let mut stats = TRACE_MAP.entry(trace).or_default();
494		stats.allocated += size as u64;
495		stats.allocations += 1;
496	}
497
498	fn take_allocation_record(&self, ptr: usize) -> Option<AllocationRecord> {
499		if self.pprof && self.effective_pprof_sample_rate() != 0 {
500			POINTER_MAP.remove(&ptr).map(|(_, record)| record)
501		} else {
502			None
503		}
504	}
505
506	fn restore_allocation_record(&self, ptr: usize, record: AllocationRecord) {
507		POINTER_MAP.insert(ptr, record);
508	}
509
510	fn finish_deallocation(&self, record: Option<AllocationRecord>, size: usize, track_stats: bool) {
511		let freed_size = record.as_ref().map(|record| record.size).unwrap_or(size);
512		if track_stats {
513			self.record_deallocation_stats(freed_size);
514		}
515
516		let Some(record) = record else {
517			return;
518		};
519
520		let mut stats = TRACE_MAP.entry(record.trace).or_default();
521		stats.freed += freed_size as u64;
522		stats.frees += 1;
523	}
524
525	#[cfg(test)]
526	fn record_deallocation(&self, ptr: usize, size: usize) {
527		let record = self.take_allocation_record(ptr);
528		self.finish_deallocation(record, size, self.stats);
529	}
530
531	#[cfg(test)]
532	fn record_reallocation(&self, old_ptr: usize, old_size: usize, new_ptr: usize, new_size: usize) {
533		let record = self.take_allocation_record(old_ptr);
534		self.finish_deallocation(record, old_size, self.stats);
535		self.record_allocation(new_ptr, new_size);
536	}
537
538	unsafe fn inner_alloc(&self, layout: Layout) -> *mut u8 {
539		match env::selected_allocator(self.default_allocator) {
540			#[cfg(feature = "allocator-jemalloc")]
541			AllocatorSelection::Jemalloc => unsafe { JEMALLOC_ALLOCATOR.alloc(layout) },
542			#[cfg(feature = "allocator-mimalloc")]
543			AllocatorSelection::Mimalloc => unsafe { MIMALLOC_ALLOCATOR.alloc(layout) },
544			_ => unsafe { System.alloc(layout) },
545		}
546	}
547
548	unsafe fn inner_alloc_zeroed(&self, layout: Layout) -> *mut u8 {
549		match env::selected_allocator(self.default_allocator) {
550			#[cfg(feature = "allocator-jemalloc")]
551			AllocatorSelection::Jemalloc => unsafe { JEMALLOC_ALLOCATOR.alloc_zeroed(layout) },
552			#[cfg(feature = "allocator-mimalloc")]
553			AllocatorSelection::Mimalloc => unsafe { MIMALLOC_ALLOCATOR.alloc_zeroed(layout) },
554			_ => unsafe { System.alloc_zeroed(layout) },
555		}
556	}
557
558	unsafe fn inner_dealloc(&self, ptr: *mut u8, layout: Layout) {
559		match env::selected_allocator(self.default_allocator) {
560			#[cfg(feature = "allocator-jemalloc")]
561			AllocatorSelection::Jemalloc => unsafe { JEMALLOC_ALLOCATOR.dealloc(ptr, layout) },
562			#[cfg(feature = "allocator-mimalloc")]
563			AllocatorSelection::Mimalloc => unsafe { MIMALLOC_ALLOCATOR.dealloc(ptr, layout) },
564			_ => unsafe { System.dealloc(ptr, layout) },
565		}
566	}
567
568	unsafe fn inner_realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
569		match env::selected_allocator(self.default_allocator) {
570			#[cfg(feature = "allocator-jemalloc")]
571			AllocatorSelection::Jemalloc => unsafe { JEMALLOC_ALLOCATOR.realloc(ptr, layout, new_size) },
572			#[cfg(feature = "allocator-mimalloc")]
573			AllocatorSelection::Mimalloc => unsafe { MIMALLOC_ALLOCATOR.realloc(ptr, layout, new_size) },
574			_ => unsafe { System.realloc(ptr, layout, new_size) },
575		}
576	}
577}
578
579fn enter_alloc<T>(func: impl FnOnce() -> T) -> T {
580	let Ok(current_value) = IN_ALLOC.try_with(|x| {
581		let current_value = x.get();
582		x.set(true);
583		current_value
584	}) else {
585		return func();
586	};
587	let output = func();
588	let _ = IN_ALLOC.try_with(|x| x.set(current_value));
589	output
590}
591
592thread_local! {
593	/// Used to avoid recursive alloc/dealloc calls for interior allocation.
594	static IN_ALLOC: Cell<bool> = const { Cell::new(false) };
595	static NEXT_SAMPLE: Cell<i64> = const { Cell::new(i64::MIN) };
596	static NEXT_SAMPLE_RATE: Cell<usize> = const { Cell::new(usize::MAX) };
597	static RNG_STATE: Cell<u64> = const { Cell::new(0) };
598	static LOCAL_STATS: LocalAllocationStats = const { LocalAllocationStats::new() };
599}
600
601static GLOBAL_STATS: stats::AtomicAllocations = stats::AtomicAllocations::new();
602static CURRENT_PPROF_SAMPLE_RATE: AtomicUsize = AtomicUsize::new(DEFAULT_PPROF_SAMPLE_RATE);
603lazy_static::lazy_static! {
604	static ref POINTER_MAP: DashMap<usize, AllocationRecord> = DashMap::new();
605	static ref TRACE_MAP: DashMap<HashedBacktrace, stats::Allocations> = DashMap::new();
606}
607
608/// Return a snapshot of coarse process-wide allocation counters.
609///
610/// These counters are updated only when the global allocator wrapper was
611/// configured with [`PprofAlloc::with_stats`]. Builds with jemalloc support
612/// report jemalloc's native current allocated bytes as `allocated`
613/// and leave cumulative free/object counters unset.
614pub fn allocation_stats() -> stats::Allocations {
615	if cfg!(feature = "allocator-jemalloc")
616		&& env::cached_allocator() == Some(AllocatorSelection::Jemalloc)
617	{
618		if let Some(stats) = native_allocation_stats() {
619			return stats;
620		}
621	}
622
623	let _ = LOCAL_STATS.try_with(|stats| stats.flush());
624	GLOBAL_STATS.snapshot()
625}
626
627#[cfg(feature = "allocator-jemalloc")]
628fn native_allocation_stats() -> Option<stats::Allocations> {
629	use tikv_jemalloc_ctl::{epoch, stats as jemalloc_stats};
630
631	epoch::advance().ok()?;
632	let allocated = jemalloc_stats::allocated::read().ok()? as u64;
633	Some(stats::Allocations {
634		allocated,
635		freed: 0,
636		allocations: 0,
637		frees: 0,
638	})
639}
640
641#[cfg(not(feature = "allocator-jemalloc"))]
642fn native_allocation_stats() -> Option<stats::Allocations> {
643	None
644}
645
646/// Return the stack capture mode compiled into this build.
647///
648/// Linux x86_64/aarch64 builds use the fast frame-pointer unwinder when the
649/// default `frame-pointer` feature is enabled. Other builds use the `backtrace`
650/// crate fallback.
651pub const fn capture_mode() -> CaptureMode {
652	trace::capture_mode()
653}
654
655fn current_pprof_sample_rate() -> usize {
656	CURRENT_PPROF_SAMPLE_RATE.load(Ordering::Relaxed)
657}
658
659fn env_pprof_sample_rate(default_rate: usize) -> usize {
660	env::pprof_sample_rate(default_rate)
661}
662
663fn should_sample_allocation(size: usize, sample_rate: usize) -> bool {
664	if size == 0 || sample_rate == 0 {
665		return false;
666	}
667	if sample_rate == 1 {
668		return true;
669	}
670
671	NEXT_SAMPLE
672		.try_with(|next_sample| {
673			NEXT_SAMPLE_RATE.try_with(|next_sample_rate| {
674				if next_sample_rate.get() != sample_rate {
675					next_sample.set(next_sample_distance(sample_rate));
676					next_sample_rate.set(sample_rate);
677				}
678
679				let next = next_sample
680					.get()
681					.saturating_sub(i64::try_from(size).unwrap_or(i64::MAX));
682				if next < 0 {
683					next_sample.set(next_sample_distance(sample_rate));
684					true
685				} else {
686					next_sample.set(next);
687					false
688				}
689			})
690		})
691		.ok()
692		.and_then(Result::ok)
693		.unwrap_or(false)
694}
695
696fn next_sample_distance(sample_rate: usize) -> i64 {
697	match sample_rate {
698		0 => i64::MAX,
699		1 => 0,
700		rate => i64::from(fast_exp_rand(rate)),
701	}
702}
703
704fn fast_exp_rand(mean: usize) -> i32 {
705	let mean = mean.min(MAX_FAST_EXP_RAND_MEAN);
706	if mean == 0 {
707		return 0;
708	}
709
710	let q = (cheap_random() % u64::from(1u32 << RANDOM_BIT_COUNT)) as u32 + 1;
711	let qlog = ((q as f64).log2() - RANDOM_BIT_COUNT as f64).min(0.0);
712	(qlog * (-std::f64::consts::LN_2 * mean as f64)) as i32 + 1
713}
714
715fn cheap_random() -> u64 {
716	RNG_STATE
717		.try_with(|state| {
718			let mut x = state.get();
719			if x == 0 {
720				x = random_seed();
721			}
722			x ^= x >> 12;
723			x ^= x << 25;
724			x ^= x >> 27;
725			state.set(x);
726			x.wrapping_mul(0x2545_f491_4f6c_dd1d)
727		})
728		.unwrap_or_else(|_| random_seed())
729}
730
731fn random_seed() -> u64 {
732	let stack_addr = &() as *const () as usize as u64;
733	let time = SystemTime::now()
734		.duration_since(UNIX_EPOCH)
735		.map(|duration| duration.as_nanos() as u64)
736		.unwrap_or(0);
737	let seed = stack_addr ^ time ^ 0x9e37_79b9_7f4a_7c15;
738	if seed == 0 {
739		0x9e37_79b9_7f4a_7c15
740	} else {
741		seed
742	}
743}
744
745/// Capture a best-effort process memory snapshot.
746///
747/// Individual probes that fail are represented as `None` in the returned
748/// [`MemorySnapshot`]. This function intentionally does not fail as a whole.
749pub fn snapshot() -> MemorySnapshot {
750	enter_alloc(|| MemorySnapshot {
751		captured_at_unix_ms: SystemTime::now()
752			.duration_since(UNIX_EPOCH)
753			.expect("system time must be after the UNIX epoch")
754			.as_millis()
755			.try_into()
756			.expect("timestamp must fit in u64"),
757		capture_mode: capture_mode(),
758		allocation_stats: allocation_stats(),
759		pprof: pprof_summary(),
760		allocator: allocator::snapshot(),
761		cgroup: stats::cgroups::get_stats().ok(),
762		smaps: stats::smaps::rollup().ok(),
763	})
764}
765
766fn pprof_summary() -> PprofSummary {
767	let mut summary = PprofSummary::default();
768	let sample_rate = current_pprof_sample_rate();
769	for entry in TRACE_MAP.iter() {
770		let stats = entry.value();
771		let values = HeapSampleValues::from_allocations(stats, sample_rate);
772		summary.total_stacks += 1;
773		summary.alloc_space_bytes += values.alloc_space.max(0) as u64;
774		summary.inuse_space_bytes += values.inuse_space.max(0) as u64;
775		summary.alloc_objects += values.alloc_objects.max(0) as u64;
776		summary.inuse_objects += values.inuse_objects.max(0) as u64;
777		if values.inuse_space > 0 {
778			summary.live_stacks += 1;
779		}
780	}
781	summary
782}
783
784#[cfg(feature = "allocator-jemalloc")]
785static JEMALLOC_ALLOCATOR: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
786#[cfg(feature = "allocator-mimalloc")]
787static MIMALLOC_ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc;
788
789unsafe impl GlobalAlloc for PprofAlloc {
790	unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
791		unsafe {
792			let tracking_mode = self.tracking_mode();
793			match tracking_mode {
794				TrackingMode::System => return System.alloc(layout),
795				#[cfg(feature = "allocator-jemalloc")]
796				TrackingMode::Jemalloc => return JEMALLOC_ALLOCATOR.alloc(layout),
797				#[cfg(feature = "allocator-mimalloc")]
798				TrackingMode::Mimalloc => return MIMALLOC_ALLOCATOR.alloc(layout),
799				_ => {},
800			}
801
802			let sample_rate = matches!(
803				tracking_mode,
804				TrackingMode::Pprof | TrackingMode::PprofStats
805			)
806			.then(|| self.effective_pprof_sample_rate());
807			let track_stats = matches!(
808				tracking_mode,
809				TrackingMode::Stats | TrackingMode::PprofStats
810			);
811			if self.tracking_is_recursive(sample_rate, track_stats) {
812				return self.inner_alloc(layout);
813			}
814
815			let ptr = self.inner_alloc(layout);
816			self.record_tracked_allocation(ptr, layout.size(), sample_rate, track_stats);
817			ptr
818		}
819	}
820
821	unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
822		unsafe {
823			let tracking_mode = self.tracking_mode();
824			match tracking_mode {
825				TrackingMode::System => return System.alloc_zeroed(layout),
826				#[cfg(feature = "allocator-jemalloc")]
827				TrackingMode::Jemalloc => return JEMALLOC_ALLOCATOR.alloc_zeroed(layout),
828				#[cfg(feature = "allocator-mimalloc")]
829				TrackingMode::Mimalloc => return MIMALLOC_ALLOCATOR.alloc_zeroed(layout),
830				_ => {},
831			}
832
833			let sample_rate = matches!(
834				tracking_mode,
835				TrackingMode::Pprof | TrackingMode::PprofStats
836			)
837			.then(|| self.effective_pprof_sample_rate());
838			let track_stats = matches!(
839				tracking_mode,
840				TrackingMode::Stats | TrackingMode::PprofStats
841			);
842			if self.tracking_is_recursive(sample_rate, track_stats) {
843				return self.inner_alloc_zeroed(layout);
844			}
845
846			let ptr = self.inner_alloc_zeroed(layout);
847			self.record_tracked_allocation(ptr, layout.size(), sample_rate, track_stats);
848			ptr
849		}
850	}
851
852	unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
853		unsafe {
854			let tracking_mode = self.tracking_mode();
855			match tracking_mode {
856				TrackingMode::System => return System.dealloc(ptr, layout),
857				#[cfg(feature = "allocator-jemalloc")]
858				TrackingMode::Jemalloc => return JEMALLOC_ALLOCATOR.dealloc(ptr, layout),
859				#[cfg(feature = "allocator-mimalloc")]
860				TrackingMode::Mimalloc => return MIMALLOC_ALLOCATOR.dealloc(ptr, layout),
861				_ => {},
862			}
863
864			let sample_rate = matches!(
865				tracking_mode,
866				TrackingMode::Pprof | TrackingMode::PprofStats
867			)
868			.then(|| self.effective_pprof_sample_rate());
869			let track_stats = matches!(
870				tracking_mode,
871				TrackingMode::Stats | TrackingMode::PprofStats
872			);
873			if self.tracking_is_recursive(sample_rate, track_stats) {
874				self.inner_dealloc(ptr, layout);
875				return;
876			}
877
878			if sample_rate.is_none() {
879				self.inner_dealloc(ptr, layout);
880				if track_stats {
881					self.record_deallocation_stats(layout.size());
882				}
883				return;
884			}
885
886			enter_alloc(|| {
887				let record = self.take_allocation_record(ptr as usize);
888				self.inner_dealloc(ptr, layout);
889				self.finish_deallocation(record, layout.size(), track_stats);
890			});
891		}
892	}
893
894	unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
895		unsafe {
896			let tracking_mode = self.tracking_mode();
897			match tracking_mode {
898				TrackingMode::System => return System.realloc(ptr, layout, new_size),
899				#[cfg(feature = "allocator-jemalloc")]
900				TrackingMode::Jemalloc => return JEMALLOC_ALLOCATOR.realloc(ptr, layout, new_size),
901				#[cfg(feature = "allocator-mimalloc")]
902				TrackingMode::Mimalloc => return MIMALLOC_ALLOCATOR.realloc(ptr, layout, new_size),
903				_ => {},
904			}
905
906			let sample_rate = matches!(
907				tracking_mode,
908				TrackingMode::Pprof | TrackingMode::PprofStats
909			)
910			.then(|| self.effective_pprof_sample_rate());
911			let track_stats = matches!(
912				tracking_mode,
913				TrackingMode::Stats | TrackingMode::PprofStats
914			);
915			if self.tracking_is_recursive(sample_rate, track_stats) {
916				return self.inner_realloc(ptr, layout, new_size);
917			}
918
919			if sample_rate.is_none() {
920				let new_ptr = self.inner_realloc(ptr, layout, new_size);
921				if !new_ptr.is_null() && track_stats {
922					self.record_deallocation_stats(layout.size());
923					self.record_allocation_stats(new_size);
924				}
925				return new_ptr;
926			}
927
928			enter_alloc(|| {
929				let record = self.take_allocation_record(ptr as usize);
930				let new_ptr = self.inner_realloc(ptr, layout, new_size);
931				if !new_ptr.is_null() {
932					self.finish_deallocation(record, layout.size(), track_stats);
933					self.record_tracked_allocation(new_ptr, new_size, sample_rate, track_stats);
934				} else if let Some(record) = record {
935					self.restore_allocation_record(ptr as usize, record);
936				}
937				new_ptr
938			})
939		}
940	}
941}
942
943pub fn generate_pprof() -> anyhow::Result<Vec<u8>> {
944	if cfg!(feature = "allocator-jemalloc")
945		&& env::cached_allocator() == Some(AllocatorSelection::Jemalloc)
946		&& env::selected_pprof_backend() == PprofBackend::Native
947	{
948		return generate_jemalloc_pprof();
949	}
950
951	enter_alloc(|| {
952		let sample_rate = current_pprof_sample_rate();
953		let mut profile = StackProfile {
954			annotations: Default::default(),
955			stacks: Default::default(),
956			mappings: if let Some(m) = crate::pprof::MAPPINGS.as_deref() {
957				m.to_vec()
958			} else {
959				Default::default()
960			},
961		};
962
963		for entry in TRACE_MAP.iter() {
964			let sample_values = HeapSampleValues::from_allocations(entry.value(), sample_rate);
965			if sample_values.alloc_space == 0
966				&& sample_values.inuse_space == 0
967				&& sample_values.alloc_objects == 0
968				&& sample_values.inuse_objects == 0
969			{
970				continue;
971			}
972
973			profile.push_stack(
974				WeightedStack {
975					addrs: entry.key().addrs(),
976					values: smallvec::smallvec![
977						sample_values.alloc_objects,
978						sample_values.alloc_space,
979						sample_values.inuse_objects,
980						sample_values.inuse_space
981					],
982				},
983				None,
984			);
985		}
986
987		Ok(profile.to_pprof_with_period(
988			&[
989				("alloc_objects", "count"),
990				("alloc_space", "bytes"),
991				("inuse_objects", "count"),
992				("inuse_space", "bytes"),
993			],
994			("space", "bytes"),
995			sample_rate as i64,
996			None,
997		))
998	})
999}
1000
1001/// Apply optional runtime configuration from environment variables.
1002///
1003/// Call this during application startup. Without `allocator-jemalloc`, this is a
1004/// no-op. With `allocator-jemalloc` and jemalloc selected, this applies
1005/// [`PPROF_BACKEND_ENV`] to jemalloc's runtime `prof.active` setting and
1006/// [`PPROF_SAMPLE_RATE_ENV`] to jemalloc's runtime `prof.lg_sample` setting.
1007/// The application is still responsible for enabling jemalloc profiling at
1008/// initialization time, for example with a `malloc_conf`/`MALLOC_CONF` that
1009/// includes `prof:true` and `prof_accum:true`.
1010#[cfg(feature = "allocator-jemalloc")]
1011pub fn configure() -> anyhow::Result<()> {
1012	configure_with_default(Allocator::System)
1013}
1014
1015/// Apply optional runtime configuration with an explicit fallback allocator.
1016///
1017/// Use this instead of [`configure`] when `PprofAlloc::with_default` is set to a
1018/// non-system allocator and startup configuration must run before the first
1019/// allocation.
1020#[cfg(feature = "allocator-jemalloc")]
1021pub fn configure_with_default(default: Allocator) -> anyhow::Result<()> {
1022	use tikv_jemalloc_ctl::raw;
1023
1024	if env::selected_allocator(default) != AllocatorSelection::Jemalloc {
1025		return Ok(());
1026	}
1027
1028	let prof_enabled: bool = unsafe { raw::read(b"opt.prof\0") }?;
1029	if !prof_enabled {
1030		anyhow::bail!(
1031			"jemalloc profiling is unavailable; configure jemalloc with prof:true before allocator initialization"
1032		);
1033	}
1034
1035	let sample_rate = env_pprof_sample_rate(DEFAULT_PPROF_SAMPLE_RATE);
1036	let active = env::selected_pprof_backend() == PprofBackend::Native && sample_rate != 0;
1037	if let Some(lg_sample) = jemalloc_lg_prof_sample(sample_rate) {
1038		unsafe { raw::write(b"prof.reset\0", lg_sample) }?;
1039	}
1040	unsafe { raw::write(b"prof.active\0", active) }?;
1041	Ok(())
1042}
1043
1044#[cfg(not(feature = "allocator-jemalloc"))]
1045pub fn configure() -> anyhow::Result<()> {
1046	Ok(())
1047}
1048
1049#[cfg(not(feature = "allocator-jemalloc"))]
1050pub fn configure_with_default(_default: Allocator) -> anyhow::Result<()> {
1051	Ok(())
1052}
1053
1054/// Generate a pprof heap profile from jemalloc's native heap profiler.
1055///
1056/// This requires the `allocator-jemalloc` feature, a jemalloc build
1057/// with profiling support, and jemalloc runtime configuration such as
1058/// `prof:true` and `prof_active:true`.
1059#[cfg(feature = "allocator-jemalloc")]
1060pub fn generate_jemalloc_pprof() -> anyhow::Result<Vec<u8>> {
1061	use std::ffi::CString;
1062	use std::fs::File;
1063	use std::io::{BufReader, Read, Seek, SeekFrom};
1064	use std::os::fd::{FromRawFd, RawFd};
1065
1066	use anyhow::Context;
1067	use tikv_jemalloc_ctl::raw;
1068
1069	let prof_enabled: bool = unsafe { raw::read(b"opt.prof\0") }?;
1070	if !prof_enabled {
1071		anyhow::bail!(
1072			"jemalloc native profiling is unavailable; enable jemalloc profiling with prof:true"
1073		);
1074	}
1075
1076	let prof_active: bool = unsafe { raw::read(b"prof.active\0") }?;
1077	if !prof_active {
1078		anyhow::bail!(
1079			"jemalloc native profiling is inactive; enable prof_active:true or activate it before dumping"
1080		);
1081	}
1082
1083	let fd = unsafe {
1084		libc::syscall(
1085			libc::SYS_memfd_create,
1086			c"pprof-alloc-jemalloc-profile".as_ptr(),
1087			libc::MFD_CLOEXEC,
1088		)
1089	};
1090	if fd < 0 {
1091		return Err(std::io::Error::last_os_error())
1092			.context("failed to create memfd for jemalloc profile dump");
1093	}
1094
1095	let mut file = unsafe { File::from_raw_fd(fd as RawFd) };
1096	let path = CString::new(format!("/proc/self/fd/{fd}"))?;
1097
1098	unsafe {
1099		raw::write(b"prof.dump\0", path.as_ptr())?;
1100	}
1101
1102	file.seek(SeekFrom::Start(0))?;
1103	let mut dump = Vec::new();
1104	file.read_to_end(&mut dump)?;
1105
1106	let mappings = crate::pprof::MAPPINGS
1107		.as_deref()
1108		.map(|mappings| mappings.to_vec())
1109		.unwrap_or_default();
1110	let parsed =
1111		crate::pprof::parse_jemalloc_heap_profile(BufReader::new(dump.as_slice()), mappings)?;
1112	Ok(parsed.profile.to_pprof_with_period(
1113		&[
1114			("alloc_objects", "count"),
1115			("alloc_space", "bytes"),
1116			("inuse_objects", "count"),
1117			("inuse_space", "bytes"),
1118		],
1119		("space", "bytes"),
1120		parsed.sampling_rate,
1121		None,
1122	))
1123}
1124
1125#[cfg(feature = "allocator-jemalloc")]
1126fn jemalloc_lg_prof_sample(sample_rate: usize) -> Option<libc::size_t> {
1127	if sample_rate == 0 {
1128		return None;
1129	}
1130
1131	let lg_sample = usize::BITS - (sample_rate - 1).leading_zeros();
1132	Some(lg_sample.min(63) as libc::size_t)
1133}
1134
1135/// Generate a pprof heap profile from jemalloc's native heap profiler.
1136#[cfg(not(feature = "allocator-jemalloc"))]
1137pub fn generate_jemalloc_pprof() -> anyhow::Result<Vec<u8>> {
1138	anyhow::bail!(
1139		"jemalloc native profiling support is not compiled in; enable the `allocator-jemalloc` feature"
1140	)
1141}
1142
1143#[doc(hidden)]
1144#[macro_export]
1145macro_rules! __pprof_alloc_register_allocator_kind {
1146	($kind:expr) => {
1147		const _: () = {
1148			#[cfg(target_os = "linux")]
1149			#[used]
1150			#[unsafe(link_section = ".init_array")]
1151			static INIT_ARRAY: extern "C" fn() = {
1152				extern "C" fn init() {
1153					$crate::allocator::configure($kind);
1154				}
1155				init
1156			};
1157		};
1158	};
1159}
1160
1161#[macro_export]
1162macro_rules! declare_allocator_kind {
1163	($kind:expr $(;)?) => {
1164		$crate::__pprof_alloc_register_allocator_kind!($kind);
1165	};
1166}
1167
1168#[cfg(test)]
1169fn reset_tracking_state() {
1170	POINTER_MAP.clear();
1171	TRACE_MAP.clear();
1172	GLOBAL_STATS.allocated.store(0, Ordering::Relaxed);
1173	GLOBAL_STATS.freed.store(0, Ordering::Relaxed);
1174	GLOBAL_STATS.allocations.store(0, Ordering::Relaxed);
1175	GLOBAL_STATS.frees.store(0, Ordering::Relaxed);
1176	let _ = LOCAL_STATS.try_with(|stats| stats.reset());
1177	CURRENT_PPROF_SAMPLE_RATE.store(1, Ordering::Relaxed);
1178	env::reset_for_tests();
1179	let _ = NEXT_SAMPLE.try_with(|next_sample| next_sample.set(i64::MIN));
1180	let _ = NEXT_SAMPLE_RATE.try_with(|next_sample_rate| next_sample_rate.set(usize::MAX));
1181}
1182
1183#[cfg(test)]
1184mod tests {
1185	use super::*;
1186	use parking_lot::Mutex;
1187
1188	static TEST_GUARD: Mutex<()> = Mutex::new(());
1189
1190	#[test]
1191	fn allocation_stats_compute_in_use_values() {
1192		let stats = stats::Allocations {
1193			allocated: 4096,
1194			freed: 1024,
1195			allocations: 4,
1196			frees: 1,
1197		};
1198
1199		assert_eq!(stats.in_use_bytes(), 3072);
1200		assert_eq!(stats.in_use_allocations(), 3);
1201	}
1202
1203	#[test]
1204	fn sample_rate_one_records_every_profile_allocation() {
1205		let _guard = TEST_GUARD.lock();
1206		reset_tracking_state();
1207
1208		let alloc = PprofAlloc::new().with_pprof_sample_rate(1);
1209		alloc.record_allocation(0x1000, 128);
1210		alloc.record_allocation(0x2000, 64);
1211
1212		assert_eq!(current_pprof_sample_rate(), 1);
1213		assert_eq!(POINTER_MAP.len(), 2);
1214		assert_eq!(pprof_summary().alloc_space_bytes, 192);
1215		assert_eq!(pprof_summary().alloc_objects, 2);
1216	}
1217
1218	#[test]
1219	fn sample_rate_zero_disables_profile_allocation_records() {
1220		let _guard = TEST_GUARD.lock();
1221		reset_tracking_state();
1222
1223		let alloc = PprofAlloc::new().with_pprof_sample_rate(0).with_stats();
1224		alloc.record_allocation(0x1000, 128);
1225		alloc.record_deallocation(0x1000, 128);
1226
1227		assert_eq!(current_pprof_sample_rate(), 0);
1228		assert!(POINTER_MAP.is_empty());
1229		assert!(TRACE_MAP.is_empty());
1230		assert_eq!(allocation_stats().allocated, 128);
1231		assert_eq!(allocation_stats().freed, 128);
1232	}
1233
1234	#[test]
1235	fn sampled_heap_values_are_scaled_to_estimates() {
1236		let (count, size) = scale_heap_sample(1, 1024, 512);
1237
1238		assert_eq!(count, 1);
1239		assert!((1180..=1190).contains(&size));
1240	}
1241
1242	#[test]
1243	#[cfg(feature = "allocator-jemalloc")]
1244	fn pprof_backend_env_can_select_wrapper_backend() {
1245		let _guard = TEST_GUARD.lock();
1246		reset_tracking_state();
1247
1248		unsafe {
1249			std::env::set_var(ALLOCATOR_ENV, "jemalloc");
1250			std::env::set_var(PPROF_BACKEND_ENV, "wrapper");
1251		}
1252		env::reset_allocator_for_tests();
1253		env::reset_pprof_backend_for_tests();
1254		assert!(!PprofAlloc::new().native_pprof_selected());
1255
1256		unsafe {
1257			std::env::set_var(PPROF_BACKEND_ENV, "pprof-alloc");
1258		}
1259		env::reset_pprof_backend_for_tests();
1260		assert!(!PprofAlloc::new().native_pprof_selected());
1261
1262		unsafe {
1263			std::env::set_var(PPROF_BACKEND_ENV, "rust");
1264		}
1265		env::reset_pprof_backend_for_tests();
1266		assert!(!PprofAlloc::new().native_pprof_selected());
1267
1268		unsafe {
1269			std::env::remove_var(PPROF_BACKEND_ENV);
1270		}
1271		env::reset_allocator_for_tests();
1272		env::reset_pprof_backend_for_tests();
1273		assert!(PprofAlloc::new().native_pprof_selected());
1274
1275		unsafe {
1276			std::env::remove_var(ALLOCATOR_ENV);
1277		}
1278		reset_tracking_state();
1279	}
1280
1281	#[test]
1282	fn allocator_compat_env_is_used_as_fallback() {
1283		let _guard = TEST_GUARD.lock();
1284		reset_tracking_state();
1285
1286		unsafe {
1287			std::env::remove_var(ALLOCATOR_ENV);
1288			std::env::set_var("ALLOCATOR", "system");
1289		}
1290
1291		assert_eq!(
1292			env::selected_allocator(Allocator::Jemalloc),
1293			AllocatorSelection::System
1294		);
1295
1296		reset_tracking_state();
1297	}
1298
1299	#[test]
1300	#[cfg(feature = "allocator-jemalloc")]
1301	fn jemalloc_sample_rate_converts_to_log2_period() {
1302		assert_eq!(jemalloc_lg_prof_sample(0), None);
1303		assert_eq!(jemalloc_lg_prof_sample(1), Some(0));
1304		assert_eq!(jemalloc_lg_prof_sample(2), Some(1));
1305		assert_eq!(jemalloc_lg_prof_sample(3), Some(2));
1306		assert_eq!(jemalloc_lg_prof_sample(DEFAULT_PPROF_SAMPLE_RATE), Some(19));
1307	}
1308
1309	#[test]
1310	fn env_sample_rate_is_read_lazily() {
1311		let _guard = TEST_GUARD.lock();
1312		reset_tracking_state();
1313
1314		unsafe {
1315			std::env::set_var(PPROF_SAMPLE_RATE_ENV, "1");
1316		}
1317		let alloc = PprofAlloc::new().with_pprof_sample_rate_from_env(DEFAULT_PPROF_SAMPLE_RATE);
1318		alloc.record_allocation(0x1000, 128);
1319		unsafe {
1320			std::env::remove_var(PPROF_SAMPLE_RATE_ENV);
1321		}
1322
1323		assert_eq!(current_pprof_sample_rate(), 1);
1324		assert_eq!(POINTER_MAP.len(), 1);
1325	}
1326
1327	#[test]
1328	fn env_sample_rate_uses_configured_default_when_unset() {
1329		let _guard = TEST_GUARD.lock();
1330		reset_tracking_state();
1331
1332		unsafe {
1333			std::env::remove_var(PPROF_SAMPLE_RATE_ENV);
1334		}
1335		let alloc = PprofAlloc::new().with_pprof_sample_rate_from_env(1);
1336		alloc.record_allocation(0x1000, 128);
1337
1338		assert_eq!(current_pprof_sample_rate(), 1);
1339		assert_eq!(POINTER_MAP.len(), 1);
1340	}
1341
1342	#[test]
1343	fn deallocation_updates_live_profile_bytes() {
1344		let _guard = TEST_GUARD.lock();
1345		reset_tracking_state();
1346
1347		let alloc = PprofAlloc::new().with_pprof().with_stats();
1348		let trace = HashedBacktrace::capture();
1349
1350		alloc.record_allocation_with_trace(0x1000, 128, trace.clone());
1351		alloc.record_allocation_with_trace(0x2000, 64, trace.clone());
1352		alloc.record_deallocation(0x1000, 128);
1353
1354		let trace_stats = TRACE_MAP.get(&trace).unwrap();
1355		assert_eq!(trace_stats.allocated, 192);
1356		assert_eq!(trace_stats.freed, 128);
1357		assert_eq!(trace_stats.allocations, 2);
1358		assert_eq!(trace_stats.frees, 1);
1359		assert_eq!(trace_stats.in_use_bytes(), 64);
1360		assert!(POINTER_MAP.contains_key(&0x2000));
1361		assert!(!POINTER_MAP.contains_key(&0x1000));
1362	}
1363
1364	#[test]
1365	fn coarse_stats_track_allocations_and_frees() {
1366		let _guard = TEST_GUARD.lock();
1367		reset_tracking_state();
1368
1369		let alloc = PprofAlloc::new().with_stats();
1370		alloc.record_allocation(0x5000, 48);
1371		alloc.record_deallocation(0x5000, 48);
1372
1373		assert_eq!(
1374			allocation_stats(),
1375			stats::Allocations {
1376				allocated: 48,
1377				freed: 48,
1378				allocations: 1,
1379				frees: 1,
1380			}
1381		);
1382	}
1383
1384	#[test]
1385	fn reallocation_updates_live_bytes_and_pointer_ownership() {
1386		let _guard = TEST_GUARD.lock();
1387		reset_tracking_state();
1388
1389		let alloc = PprofAlloc::new().with_pprof_sample_rate(1);
1390		alloc.record_allocation_with_trace(0x3000, 32, HashedBacktrace::capture());
1391		alloc.record_reallocation(0x3000, 32, 0x4000, 96);
1392
1393		let total_live_bytes: u64 = TRACE_MAP
1394			.iter()
1395			.map(|entry| entry.value().in_use_bytes())
1396			.sum();
1397		assert_eq!(total_live_bytes, 96);
1398		assert_eq!(POINTER_MAP.get(&0x4000).unwrap().size, 96);
1399		assert!(!POINTER_MAP.contains_key(&0x3000));
1400	}
1401
1402	#[test]
1403	fn snapshot_reports_current_pprof_summary() {
1404		let _guard = TEST_GUARD.lock();
1405		reset_tracking_state();
1406
1407		let alloc = PprofAlloc::new().with_pprof();
1408		let trace = HashedBacktrace::capture();
1409
1410		alloc.record_allocation_with_trace(0x1000, 128, trace.clone());
1411		alloc.record_allocation_with_trace(0x2000, 64, trace);
1412		alloc.record_deallocation(0x1000, 128);
1413
1414		let snapshot = snapshot();
1415		assert_eq!(snapshot.capture_mode, capture_mode());
1416		assert_eq!(
1417			snapshot.pprof,
1418			PprofSummary {
1419				total_stacks: 1,
1420				live_stacks: 1,
1421				alloc_space_bytes: 192,
1422				inuse_space_bytes: 64,
1423				alloc_objects: 2,
1424				inuse_objects: 1,
1425			}
1426		);
1427	}
1428}