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
149static STACKS: SpinLock<Stacks> = SpinLock::new(Stacks {
150    table: null_mut(),
151    cap: 0,
152    len: 0,
153    all: null_mut(),
154});
155
156impl Stacks {
157    unsafe fn find_or_insert(&mut self, hash: u64, frames: &[usize]) -> *mut Stack {
158        if (self.len + 1) * 2 > self.cap && !self.grow() {
159            return null_mut();
160        }
161        let mask = self.cap - 1;
162        let mut i = hash as usize & mask;
163        loop {
164            let s = *self.table.add(i);
165            if s.is_null() {
166                break;
167            }
168            if (*s).hash == hash && Stack::frames(s) == frames {
169                return s;
170            }
171            i = (i + 1) & mask;
172        }
173        let bytes = size_of::<Stack>() + std::mem::size_of_val(frames);
174        let s: *mut Stack = base::alloc(bytes, align_of::<Stack>()).cast();
175        if s.is_null() {
176            return s;
177        }
178        ptr::write(
179            s,
180            Stack {
181                next: self.all,
182                hash,
183                depth: frames.len(),
184                alloc_objects: 0,
185                alloc_bytes: 0,
186                inuse_objects: 0,
187                inuse_bytes: 0,
188            },
189        );
190        ptr::copy_nonoverlapping(frames.as_ptr(), s.add(1).cast::<usize>(), frames.len());
191        self.all = s;
192        *self.table.add(i) = s;
193        self.len += 1;
194        s
195    }
196
197    unsafe fn grow(&mut self) -> bool {
198        let cap = (self.cap * 2).max(1024);
199        let bytes = os::round_up(cap * size_of::<usize>(), os::page_size());
200        let table: *mut *mut Stack = os::map(bytes).cast();
201        if table.is_null() {
202            return false;
203        }
204        let mut s = self.all;
205        while !s.is_null() {
206            let mut i = (*s).hash as usize & (cap - 1);
207            while !(*table.add(i)).is_null() {
208                i = (i + 1) & (cap - 1);
209            }
210            *table.add(i) = s;
211            s = (*s).next;
212        }
213        if !self.table.is_null() {
214            let old = os::round_up(self.cap * size_of::<usize>(), os::page_size());
215            os::unmap(self.table.cast(), old);
216            crate::stats::METADATA_BYTES.sub(old);
217        }
218        crate::stats::METADATA_BYTES.add(bytes);
219        self.table = table;
220        self.cap = cap;
221        true
222    }
223}
224
225#[cfg(all(unix, not(miri)))]
226/// Holds the profiler's locks across `fork`, stacks before samples
227pub(crate) fn fork_lock() {
228    STACKS.acquire();
229    SAMPLES.fork_lock();
230}
231
232#[cfg(all(unix, not(miri)))]
233pub(crate) unsafe fn fork_unlock() {
234    SAMPLES.fork_unlock();
235    STACKS.release();
236}
237
238fn hash_frames(frames: &[usize]) -> u64 {
239    let mut h: u64 = 0;
240    for &f in frames {
241        h = (h.rotate_left(5) ^ f as u64).wrapping_mul(0x517c_c1b7_2722_0a95);
242    }
243    h
244}
245
246/// Allocates and records a sample; called when a thread's sampling
247/// counter runs out. `usable` is what the allocation took off the counter
248/// and `interval` the mean of the wait that ran out, which together set
249/// the chance it had to be sampled.
250#[inline(never)]
251pub(crate) unsafe fn sample_alloc(
252    size: usize,
253    align: usize,
254    zero: bool,
255    usable: usize,
256    interval: usize,
257) -> *mut u8 {
258    // A free of this allocation on any thread is ordered after this store,
259    // so it sees that frees must check for samples
260    EVER.store(true, Ordering::Relaxed);
261    let mut frames = [0usize; MAX_FRAMES];
262    // Leave out this function
263    let depth = backtrace::capture(&mut frames, 1);
264    let frames = &frames[..depth];
265
266    let sample = SAMPLES.alloc();
267    if sample.is_null() {
268        return null_mut();
269    }
270    let ptr = alloc_sampled(size, align, zero, sample);
271    if ptr.is_null() {
272        SAMPLES.free(sample);
273        return ptr;
274    }
275
276    // Each sample stands for 1 / P(sampled) allocations of its size
277    let usable = usable.max(1) as f64;
278    let scale = 1.0 / (1.0 - (-usable / interval as f64).exp());
279    let one = f64::from(1u32 << WEIGHT_SHIFT);
280    let objects = (scale * one).round() as u64;
281    let bytes = (usable * scale * one).round() as u64;
282    let hash = hash_frames(frames);
283    let stack = {
284        let mut stacks = STACKS.lock();
285        let s = stacks.find_or_insert(hash, frames);
286        if !s.is_null() {
287            (*s).alloc_objects += u128::from(objects);
288            (*s).alloc_bytes += u128::from(bytes);
289            (*s).inuse_objects += u128::from(objects);
290            (*s).inuse_bytes += u128::from(bytes);
291        }
292        s
293    };
294    ptr::write(
295        sample,
296        Sample {
297            stack,
298            objects,
299            bytes,
300        },
301    );
302    ptr
303}
304
305/// Memory for a sample: a page run in the sampled arena, or a huge mapping
306unsafe fn alloc_sampled(size: usize, align: usize, zero: bool, sample: *mut Sample) -> *mut u8 {
307    let align = align.max(1);
308    let pages = (os::round_up(size.max(1), PAGE) + align.max(PAGE) - PAGE) >> PAGE_SHIFT;
309    if pages << PAGE_SHIFT > LARGE_MAX {
310        return huge::alloc(size, align, zero, sample);
311    }
312    let arena = arena::sampled();
313    if arena.is_null() {
314        return null_mut();
315    }
316    let run = (*arena).alloc_pages(pages, u8::MAX);
317    if run.is_null() {
318        return null_mut();
319    }
320    (*run).sample = sample;
321    let p = run_addr(run).map_addr(|a| os::round_up(a, align));
322    if zero {
323        ptr::write_bytes(p, 0, size);
324    }
325    p
326}
327
328/// Frees a sampled allocation and takes it out of the live counts
329pub(crate) unsafe fn sample_free(ptr: *mut u8) {
330    let huge = (*head_of(ptr)).kind == KIND_HUGE;
331    let sample = if huge {
332        (*huge::head(ptr)).sample
333    } else {
334        (*arena::sampled_run(ptr)).sample
335    };
336    let s = &*sample;
337    if !s.stack.is_null() {
338        let _stacks = STACKS.lock();
339        (*s.stack).inuse_objects -= u128::from(s.objects);
340        (*s.stack).inuse_bytes -= u128::from(s.bytes);
341    }
342    SAMPLES.free(sample);
343    if huge {
344        huge::free(ptr);
345    } else {
346        arena::free_sampled_run(ptr);
347    }
348}
349
350/// Counts of one stack, copied out for a dump
351struct StackCounts {
352    frames: &'static [usize],
353    alloc_objects: u64,
354    alloc_bytes: u64,
355    inuse_objects: u64,
356    inuse_bytes: u64,
357}
358
359/// Copies every stack's counts without allocating under the lock
360fn snapshot() -> Vec<StackCounts> {
361    let mut out = Vec::new();
362    loop {
363        let need = STACKS.lock().len;
364        out.reserve(need + 64);
365        let stacks = STACKS.lock();
366        if stacks.len > out.capacity() {
367            continue;
368        }
369        let round = |w: u128| ((w + (1 << (WEIGHT_SHIFT - 1))) >> WEIGHT_SHIFT) as u64;
370        let mut s = stacks.all;
371        while !s.is_null() {
372            unsafe {
373                out.push(StackCounts {
374                    frames: Stack::frames(s),
375                    alloc_objects: round((*s).alloc_objects),
376                    alloc_bytes: round((*s).alloc_bytes),
377                    inuse_objects: round((*s).inuse_objects),
378                    inuse_bytes: round((*s).inuse_bytes),
379                });
380                s = (*s).next;
381            }
382        }
383        return out;
384    }
385}
386
387/// The heap profile in pprof's protobuf format, uncompressed (`pprof`
388/// reads both plain and gzipped files). Its sample types follow Go's heap
389/// profiles: `alloc_objects`, `alloc_space`, `inuse_objects` and
390/// `inuse_space`, the default.
391///
392/// # Errors
393///
394/// [`DumpError::NotActivated`] when profiling was never activated.
395pub fn dump_pprof() -> Result<Vec<u8>, DumpError> {
396    if !EVER.load(Ordering::Relaxed) {
397        return Err(DumpError::NotActivated);
398    }
399    Ok(pprof::encode(&snapshot(), sample_interval()))
400}
401
402/// Writes the heap profile to `path`; returns its size in bytes
403///
404/// # Errors
405///
406/// [`DumpError::NotActivated`] when profiling was never activated, and
407/// [`DumpError::Io`] when the file cannot be written.
408pub fn write_pprof(path: impl AsRef<std::path::Path>) -> Result<usize, DumpError> {
409    let profile = dump_pprof()?;
410    std::fs::write(path, &profile).map_err(DumpError::Io)?;
411    Ok(profile.len())
412}
stoolap_jemalloc/prof/mod.rs

stoolap_jemalloc/prof/
mod.rs
