1pub 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
33pub 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
72pub struct PprofAlloc {
77 pprof: bool,
79 stats: bool,
81 pprof_sample_rate: usize,
83 pprof_sample_rate_from_env: bool,
85 resolved_pprof_sample_rate: AtomicUsize,
87 tracking_mode: AtomicU8,
89 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)]
230pub struct PprofSummary {
235 pub total_stacks: u64,
237 pub live_stacks: u64,
239 pub alloc_space_bytes: u64,
241 pub inuse_space_bytes: u64,
243 pub alloc_objects: u64,
245 pub inuse_objects: u64,
247}
248
249#[derive(Clone, Debug, Serialize)]
250pub struct MemorySnapshot {
255 pub captured_at_unix_ms: u64,
257 pub capture_mode: CaptureMode,
259 pub allocation_stats: stats::Allocations,
261 pub pprof: PprofSummary,
263 pub allocator: allocator::AllocatorSnapshot,
265 pub cgroup: Option<stats::cgroups::MemoryStat>,
267 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 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 pub const fn with_default(mut self, allocator: Allocator) -> Self {
296 self.default_allocator = allocator;
297 self
298 }
299
300 pub const fn with_pprof(mut self) -> Self {
306 self.pprof = true;
307 self
308 }
309
310 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 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 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 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
608pub 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
646pub 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
745pub 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#[cfg(feature = "allocator-jemalloc")]
1011pub fn configure() -> anyhow::Result<()> {
1012 configure_with_default(Allocator::System)
1013}
1014
1015#[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#[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#[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}