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

stoolap_jemalloc/prof/
mod.rs
