stoolap_jemalloc/prof/
mod.rs

1// Copyright 2026 Stoolap Contributors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Sampling heap profiler.
16//!
17//! As in jemalloc, a thread takes a sample after a random number of
18//! allocated bytes, exponentially distributed around the sample interval,
19//! so each allocation is sampled with probability `1 - exp(-size /
20//! interval)`. A sample records the allocation's stack; counts per stack
21//! are scaled by the inverse of that probability, which makes them
22//! unbiased estimates of the real totals.
23//!
24//! Sampled allocations live in chunks of their own, so a free only needs
25//! to look at the chunk header to know whether it frees a sample.
26
27mod backtrace;
28mod maps;
29mod pprof;
30
31use crate::arena::{self, run_addr};
32use crate::base::{self, Pool};
33use crate::chunk::{KIND_HUGE, head_of};
34use crate::huge;
35use crate::lock::SpinLock;
36use crate::os;
37use crate::size_class::{LARGE_MAX, PAGE, PAGE_SHIFT};
38use core::ptr::{self, null_mut};
39use core::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
40
41pub use pprof::DumpError;
42
43/// Set once profiling was ever activated; frees check for samples from then on
44pub(crate) static EVER: AtomicBool = AtomicBool::new(false);
45pub(crate) static ACTIVE: AtomicBool = AtomicBool::new(false);
46static INTERVAL: AtomicUsize = AtomicUsize::new(DEFAULT_INTERVAL);
47
48/// 512 KiB, jemalloc's default (`lg_prof_sample:19`)
49pub const DEFAULT_INTERVAL: usize = 1 << 19;
50const MAX_FRAMES: usize = 128;
51
52/// Starts sampling allocations
53pub fn activate() {
54    EVER.store(true, Ordering::SeqCst);
55    ACTIVE.store(true, Ordering::SeqCst);
56}
57
58/// Stops taking new samples; live samples stay in the profile
59pub fn deactivate() {
60    ACTIVE.store(false, Ordering::SeqCst);
61}
62
63/// Whether allocations are being sampled
64pub fn is_active() -> bool {
65    ACTIVE.load(Ordering::Relaxed)
66}
67
68/// Sets the mean number of bytes between samples, from 1 byte to 128 TiB;
69/// values outside are clamped
70pub fn set_sample_interval(bytes: usize) {
71    INTERVAL.store(bytes.clamp(1, MAX_INTERVAL), Ordering::Relaxed);
72}
73
74/// The largest interval: a sample's weight in bytes, about the interval in
75/// fixed point, then still fits in a `u64`
76const MAX_INTERVAL: usize = {
77    let max = 1u64 << 47;
78    if max < usize::MAX as u64 {
79        max as usize
80    } else {
81        usize::MAX
82    }
83};
84
85/// The mean number of bytes between samples
86pub fn sample_interval() -> usize {
87    INTERVAL.load(Ordering::Relaxed)
88}
89
90/// Weights are fixed point with this many fractional bits, so that
91/// fractional estimates add up exactly and are rounded only in a dump
92const WEIGHT_SHIFT: u32 = 16;
93/// Waits never exceed this, so that the event counters cannot overflow
94pub(crate) const MAX_WAIT: i64 = 1 << 62;
95
96/// A thread's random number seed: each thread gets its own, also when it
97/// reuses the cache memory of a thread that exited
98pub(crate) fn new_seed() -> u64 {
99    static NEXT: AtomicU64 = AtomicU64::new(0);
100    let n = NEXT.fetch_add(0x9e37_79b9_7f4a_7c15, Ordering::Relaxed)
101        ^ (u64::from(std::process::id()) << 32);
102    // splitmix64's finalizer; xorshift needs a non-zero state
103    let mut z = n;
104    z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
105    z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
106    (z ^ (z >> 31)) | 1
107}
108
109/// Bytes until the next sample: exponentially distributed with `interval`
110/// as mean
111pub(crate) fn next_wait(rng: &mut u64, interval: usize) -> i64 {
112    // xorshift64*
113    *rng ^= *rng >> 12;
114    *rng ^= *rng << 25;
115    *rng ^= *rng >> 27;
116    let r = rng.wrapping_mul(0x2545_f491_4f6c_dd1d);
117    let u = ((r >> 11) + 1) as f64 / (1u64 << 53) as f64;
118    let wait = -u.ln() * interval as f64;
119    (wait as i64).clamp(1, MAX_WAIT)
120}
121
122/// A live sampled allocation, with the fixed point weights it added
123pub struct Sample {
124    stack: *mut Stack,
125    objects: u64,
126    bytes: u64,
127}
128
129static SAMPLES: Pool<Sample> = Pool::new();
130
131/// A unique stack with its counts: estimates in fixed point
132#[repr(C)]
133struct Stack {
134    next: *mut Stack,
135    hash: u64,
136    depth: usize,
137    alloc_objects: u128,
138    alloc_bytes: u128,
139    inuse_objects: u128,
140    inuse_bytes: u128,
141    // `depth` frames follow
142}
143
144impl Stack {
145    /// The frames stored after the stack; they never change
146    unsafe fn frames<'a>(s: *const Stack) -> &'a [usize] {
147        core::slice::from_raw_parts(s.add(1).cast::<usize>(), (*s).depth)
148    }
149}
150
151/// Stacks by hash, in an open addressing table that grows by doubling.
152/// Stacks never move or go away, so a dump can read their frames
153/// without the lock.
154struct Stacks {
155    table: *mut *mut Stack,
156    cap: usize,
157    len: usize,
158    all: *mut Stack,
159}
160
161/// Held through a dump, and across `fork` before every other lock: no
162/// dump of this crate's then holds the symbolizer's lock in the child.
163/// Other users of the `backtrace` crate can still hold it at a fork.
164static DUMP: SpinLock<()> = SpinLock::new(());
165
166static STACKS: SpinLock<Stacks> = SpinLock::new(Stacks {
167    table: null_mut(),
168    cap: 0,
169    len: 0,
170    all: null_mut(),
171});
172
173impl Stacks {
174    unsafe fn find_or_insert(&mut self, hash: u64, frames: &[usize]) -> *mut Stack {
175        if (self.len + 1) * 2 > self.cap && !self.grow() {
176            return null_mut();
177        }
178        let mask = self.cap - 1;
179        let mut i = hash as usize & mask;
180        loop {
181            let s = *self.table.add(i);
182            if s.is_null() {
183                break;
184            }
185            if (*s).hash == hash && Stack::frames(s) == frames {
186                return s;
187            }
188            i = (i + 1) & mask;
189        }
190        let bytes = size_of::<Stack>() + std::mem::size_of_val(frames);
191        let s: *mut Stack = base::alloc(bytes, align_of::<Stack>()).cast();
192        if s.is_null() {
193            return s;
194        }
195        ptr::write(
196            s,
197            Stack {
198                next: self.all,
199                hash,
200                depth: frames.len(),
201                alloc_objects: 0,
202                alloc_bytes: 0,
203                inuse_objects: 0,
204                inuse_bytes: 0,
205            },
206        );
207        ptr::copy_nonoverlapping(frames.as_ptr(), s.add(1).cast::<usize>(), frames.len());
208        self.all = s;
209        *self.table.add(i) = s;
210        self.len += 1;
211        s
212    }
213
214    unsafe fn grow(&mut self) -> bool {
215        let cap = (self.cap * 2).max(1024);
216        let bytes = os::round_up(cap * size_of::<usize>(), os::page_size());
217        let table: *mut *mut Stack = os::map(bytes).cast();
218        if table.is_null() {
219            return false;
220        }
221        let mut s = self.all;
222        while !s.is_null() {
223            let mut i = (*s).hash as usize & (cap - 1);
224            while !(*table.add(i)).is_null() {
225                i = (i + 1) & (cap - 1);
226            }
227            *table.add(i) = s;
228            s = (*s).next;
229        }
230        if !self.table.is_null() {
231            let old = os::round_up(self.cap * size_of::<usize>(), os::page_size());
232            os::unmap(self.table.cast(), old);
233            crate::stats::METADATA_BYTES.sub(old);
234        }
235        crate::stats::METADATA_BYTES.add(bytes);
236        self.table = table;
237        self.cap = cap;
238        true
239    }
240}
241
242#[cfg(all(unix, not(miri)))]
243/// Waits for a dump to end and holds off the next until after `fork`
244pub(crate) fn fork_lock_dump() {
245    DUMP.acquire();
246}
247
248#[cfg(all(unix, not(miri)))]
249pub(crate) unsafe fn fork_unlock_dump() {
250    DUMP.release();
251}
252
253#[cfg(all(unix, not(miri)))]
254/// Holds the profiler's locks across `fork`, stacks before samples
255pub(crate) fn fork_lock() {
256    STACKS.acquire();
257    SAMPLES.fork_lock();
258}
259
260#[cfg(all(unix, not(miri)))]
261pub(crate) unsafe fn fork_unlock() {
262    SAMPLES.fork_unlock();
263    STACKS.release();
264}
265
266fn hash_frames(frames: &[usize]) -> u64 {
267    let mut h: u64 = 0;
268    for &f in frames {
269        h = (h.rotate_left(5) ^ f as u64).wrapping_mul(0x517c_c1b7_2722_0a95);
270    }
271    h
272}
273
274/// Allocates and records a sample; called when a thread's sampling
275/// counter runs out. `usable` is what the allocation took off the counter
276/// and `interval` the mean of the wait that ran out, which together set
277/// the chance it had to be sampled.
278#[inline(never)]
279pub(crate) unsafe fn sample_alloc(
280    size: usize,
281    align: usize,
282    zero: bool,
283    usable: usize,
284    interval: usize,
285) -> *mut u8 {
286    // A free of this allocation on any thread is ordered after this store,
287    // so it sees that frees must check for samples
288    EVER.store(true, Ordering::Relaxed);
289    let mut frames = [0usize; MAX_FRAMES];
290    // Leave out this function
291    let depth = backtrace::capture(&mut frames, 1);
292    let frames = &frames[..depth];
293
294    let sample = SAMPLES.alloc();
295    if sample.is_null() {
296        return null_mut();
297    }
298    let ptr = alloc_sampled(size, align, zero, sample);
299    if ptr.is_null() {
300        SAMPLES.free(sample);
301        return ptr;
302    }
303
304    // Each sample stands for 1 / P(sampled) allocations of its size
305    let usable = usable.max(1) as f64;
306    // 1 - e^-x, exact also when x is tiny
307    let scale = 1.0 / -(-usable / interval as f64).exp_m1();
308    let one = f64::from(1u32 << WEIGHT_SHIFT);
309    let objects = (scale * one).round() as u64;
310    let bytes = (usable * scale * one).round() as u64;
311    let hash = hash_frames(frames);
312    let stack = {
313        let mut stacks = STACKS.lock();
314        let s = stacks.find_or_insert(hash, frames);
315        if !s.is_null() {
316            (*s).alloc_objects += u128::from(objects);
317            (*s).alloc_bytes += u128::from(bytes);
318            (*s).inuse_objects += u128::from(objects);
319            (*s).inuse_bytes += u128::from(bytes);
320        }
321        s
322    };
323    ptr::write(
324        sample,
325        Sample {
326            stack,
327            objects,
328            bytes,
329        },
330    );
331    ptr
332}
333
334/// Memory for a sample: a page run in the sampled arena, or a huge mapping
335unsafe fn alloc_sampled(size: usize, align: usize, zero: bool, sample: *mut Sample) -> *mut u8 {
336    let align = align.max(1);
337    let pages = (os::round_up(size.max(1), PAGE) + align.max(PAGE) - PAGE) >> PAGE_SHIFT;
338    if pages << PAGE_SHIFT > LARGE_MAX {
339        return huge::alloc(size, align, zero, sample);
340    }
341    let arena = arena::sampled();
342    if arena.is_null() {
343        return null_mut();
344    }
345    let run = (*arena).alloc_pages(pages, u8::MAX);
346    if run.is_null() {
347        return null_mut();
348    }
349    (*run).sample = sample;
350    let p = run_addr(run).map_addr(|a| os::round_up(a, align));
351    if zero {
352        ptr::write_bytes(p, 0, size);
353    }
354    p
355}
356
357/// Frees a sampled allocation and takes it out of the live counts
358pub(crate) unsafe fn sample_free(ptr: *mut u8) {
359    let huge = (*head_of(ptr)).kind == KIND_HUGE;
360    let sample = if huge {
361        (*huge::head(ptr)).sample
362    } else {
363        (*arena::sampled_run(ptr)).sample
364    };
365    let s = &*sample;
366    if !s.stack.is_null() {
367        let _stacks = STACKS.lock();
368        (*s.stack).inuse_objects -= u128::from(s.objects);
369        (*s.stack).inuse_bytes -= u128::from(s.bytes);
370    }
371    SAMPLES.free(sample);
372    if huge {
373        huge::free(ptr);
374    } else {
375        arena::free_sampled_run(ptr);
376    }
377}
378
379/// Counts of one stack, copied out for a dump
380struct StackCounts {
381    frames: &'static [usize],
382    alloc_objects: u64,
383    alloc_bytes: u64,
384    inuse_objects: u64,
385    inuse_bytes: u64,
386}
387
388/// Copies every stack's counts without allocating under the lock
389fn snapshot() -> Vec<StackCounts> {
390    let mut out = Vec::new();
391    loop {
392        let need = STACKS.lock().len;
393        out.reserve(need + 64);
394        let stacks = STACKS.lock();
395        if stacks.len > out.capacity() {
396            continue;
397        }
398        let round = |w: u128| ((w + (1 << (WEIGHT_SHIFT - 1))) >> WEIGHT_SHIFT) as u64;
399        let mut s = stacks.all;
400        while !s.is_null() {
401            unsafe {
402                out.push(StackCounts {
403                    frames: Stack::frames(s),
404                    alloc_objects: round((*s).alloc_objects),
405                    alloc_bytes: round((*s).alloc_bytes),
406                    inuse_objects: round((*s).inuse_objects),
407                    inuse_bytes: round((*s).inuse_bytes),
408                });
409                s = (*s).next;
410            }
411        }
412        return out;
413    }
414}
415
416/// The heap profile in pprof's protobuf format, uncompressed (`pprof`
417/// reads both plain and gzipped files). Its sample types follow Go's heap
418/// profiles: `alloc_objects`, `alloc_space`, `inuse_objects` and
419/// `inuse_space`, the default.
420///
421/// # Errors
422///
423/// [`DumpError::NotActivated`] when profiling was never activated.
424pub fn dump_pprof() -> Result<Vec<u8>, DumpError> {
425    if !EVER.load(Ordering::Relaxed) {
426        return Err(DumpError::NotActivated);
427    }
428    let _dump = DUMP.lock();
429    let profile = pprof::encode(&snapshot(), sample_interval());
430    // The symbolizer keeps the debug information it parsed, tens of
431    // megabytes for a large library
432    #[cfg(feature = "symbolize")]
433    ::backtrace::clear_symbol_cache();
434    Ok(profile)
435}
436
437/// Writes the heap profile to `path`; returns its size in bytes
438///
439/// # Errors
440///
441/// [`DumpError::NotActivated`] when profiling was never activated, and
442/// [`DumpError::Io`] when the file cannot be written.
443pub fn write_pprof(path: impl AsRef<std::path::Path>) -> Result<usize, DumpError> {
444    let profile = dump_pprof()?;
445    std::fs::write(path, &profile).map_err(DumpError::Io)?;
446    Ok(profile.len())
447}
stoolap_jemalloc/prof/mod.rs

stoolap_jemalloc/prof/
mod.rs
