diff --git a/CHANGELOG.md b/CHANGELOG.md index b185c866c..53f13b88e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,12 @@ +# Unreleased + +- jemalloc-ctl: expose `prof.active`, `prof.lg_sample`, and `prof.reset` via + `profiling::{prof_active, lg_sample, prof_reset}` +- jemalloc-ctl: expose jemalloc's experimental sample hooks under the + `profiling` feature (`set_prof_sample_hook`, `set_prof_sample_free_hook`, + `set_prof_backtrace_hook`, `noop_prof_backtrace_hook`, and the `Prof*Hook` + types) + # 0.7.0 - 2026-05-25 - Reverse order of MAKEFLAGS priority (#152) diff --git a/ci/run.sh b/ci/run.sh index 56ef68600..621995cdf 100755 --- a/ci/run.sh +++ b/ci/run.sh @@ -59,10 +59,15 @@ case "${TARGET}" in cargo test --target "${TARGET}" \ --manifest-path jemalloc-ctl/Cargo.toml \ --no-default-features - # FIXME: cross fails to pass features to jemalloc-ctl - # ${CARGO_CMD} test --target "${TARGET}" \ - # --manifest-path jemalloc-ctl \ - # --no-default-features --features use_std + ( + # This malloc_conf is for the Rust hook tests and breaks jemalloc's C tests. + unset JEMALLOC_SYS_RUN_JEMALLOC_TESTS + JEMALLOC_SYS_WITH_MALLOC_CONF=prof:true,prof_active:false \ + cargo test --target "${TARGET}" \ + --manifest-path jemalloc-ctl/Cargo.toml \ + --no-default-features \ + --features 'profiling use_std' + ) ;; esac diff --git a/jemalloc-ctl/src/macros.rs b/jemalloc-ctl/src/macros.rs index 94050e354..85b6169f0 100644 --- a/jemalloc-ctl/src/macros.rs +++ b/jemalloc-ctl/src/macros.rs @@ -130,6 +130,8 @@ macro_rules! make_test { "background_thread" | "max_background_threads" if cfg!(target_os = "macos") => return, + // Requires `opt.prof` and can race with hook tests. + "prof_active" if cfg!(feature = "profiling") => return, _ => (), } diff --git a/jemalloc-ctl/src/profiling.rs b/jemalloc-ctl/src/profiling.rs index 306ab1a8e..a77e52e45 100644 --- a/jemalloc-ctl/src/profiling.rs +++ b/jemalloc-ctl/src/profiling.rs @@ -1,6 +1,21 @@ //! `jemalloc`'s run-time configuration for profiling-specific settings. //! //! These settings are controlled by the `MALLOC_CONF` environment variable. +//! +//! This module also exposes on-the-fly control via [`prof_active`] and +//! [`prof_reset`], along with `jemalloc`'s experimental +//! `experimental.hooks.prof_sample`/`prof_sample_free`/`prof_backtrace` hooks +//! via [`set_prof_sample_hook`], [`set_prof_sample_free_hook`], and +//! [`set_prof_backtrace_hook`]. +//! +//! # Experimental hook API +//! +//! `jemalloc` considers these hook mallctls experimental. Their names and +//! callback ABIs may change between versions without notice. Hooks are +//! process-wide, may run concurrently, and replacement does not wait for +//! in-flight calls. + +use libc::{c_uint, c_void}; option! { lg_prof_interval[ str: b"opt.lg_prof_interval\0", non_str: 2 ] => libc::ssize_t | @@ -153,3 +168,330 @@ option! { /// ``` mib_docs: /// See [`prof_leak`]. } + +option! { + prof_active[ str: b"prof.active\0", non_str: 2 ] => bool | + ops: r,w,u | + docs: + /// On-the-fly activation/deactivation of memory profiling. + /// + /// This is a secondary control mechanism on top of `opt.prof`, and is + /// only effective once `opt.prof` is `true`. When it is `false`, reading + /// [`prof_active`] returns `false`; writing `false` is accepted as a no-op, + /// while writing `true` fails with `ENOENT`. `jemalloc` initialises + /// [`prof_active`] to + /// `opt.prof_active` (which itself defaults to `true`) as soon as + /// `opt.prof` is `true`, so a configuration with `opt.prof` enabled samples + /// by default unless [`prof_active`] is set to `false`, e.g. via + /// `prof_active:false` in `MALLOC_CONF`. + /// + /// `thread.prof.active` is a separate per-thread gate, initialised from + /// `opt.prof_thread_active_init`; both it and this global control must be + /// active for a thread to sample. + /// + /// While this control is inactive, no new allocation samples are selected, + /// so [`ProfSampleHook`] does not fire. [`ProfSampleFreeHook`] can still + /// fire for allocations sampled before deactivation. + /// + /// # Examples + /// + /// ``` + /// # #[global_allocator] + /// # static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + /// # + /// # fn main() { + /// use tikv_jemalloc_ctl::profiling; + /// // `false` is always accepted, even if `opt.prof` is disabled at + /// // runtime; writing `true` additionally requires `opt.prof` to be + /// // `true`, else it fails with `ENOENT`. + /// let was_active = profiling::prof_active::update(false).unwrap(); + /// profiling::prof_active::write(was_active).unwrap(); + /// # } + /// ``` + mib_docs: /// See [`prof_active`]. +} + +option! { + lg_sample[ str: b"prof.lg_sample\0", non_str: 2 ] => libc::size_t | + ops: r | + docs: + /// Current log base 2 of the mean number of bytes between samples. + /// + /// Initialised from [`lg_prof_sample`] and updated by [`prof_reset`]. + mib_docs: /// See [`lg_sample`]. +} + +/// Resets `jemalloc`'s heap profile sample accumulators and, going forward, +/// draws sample intervals from a geometric distribution with a mean of +/// `2^new_lg_sample` bytes of allocation activity (values `>= 64` are clamped +/// to `63`). Sampling still requires `prof.active` and `thread.prof.active`. +/// +/// Corresponds to `prof.reset`, which is write-only: unlike most keys in +/// this module, there is no matching `read()`/`update()`. +/// +/// # Errors +/// +/// Returns an error (`ENOENT`) if `opt.prof` is `false` at runtime. The +/// `profiling` feature enables profiling support but does not set `opt.prof`; +/// configure `prof:true`, for example via `JEMALLOC_SYS_WITH_MALLOC_CONF` at +/// build time. +pub fn prof_reset(new_lg_sample: libc::size_t) -> crate::error::Result<()> { + unsafe { crate::raw::write(b"prof.reset\0", new_lg_sample) } +} + +/// Signature of a hook installable via [`set_prof_sample_hook`]. +/// +/// `jemalloc` invokes this hook synchronously, inline on the allocating +/// thread, immediately after it decides to sample an allocation of +/// `usable_size` bytes at `ptr` (the request was for `size` bytes; +/// `usable_size` is jemalloc's usable/rounded-up size). `backtrace` points +/// to an array of `backtrace_length` `void *` frames captured by the installed +/// [`ProfBacktraceHook`] (`backtrace_length` is `0` if +/// [`noop_prof_backtrace_hook`] is installed). +/// +/// # Safety +/// +/// Arguments are valid only during the call. The hook must not deallocate +/// `ptr`, retain any argument, or mutate `backtrace`. No `jemalloc` mutex is +/// held, and the hook may allocate or free other allocations; nested allocator +/// activity is excluded from profiling. Hooks may run concurrently and must +/// not unwind across the `extern "C"` boundary. +pub type ProfSampleHook = unsafe extern "C" fn( + ptr: *const c_void, + size: libc::size_t, + backtrace: *mut *mut c_void, + backtrace_length: c_uint, + usable_size: libc::size_t, +); + +/// Signature of a hook installable via [`set_prof_sample_free_hook`]. +/// +/// `jemalloc` invokes this hook synchronously, inline on the freeing +/// thread, just before it frees a previously-sampled allocation of +/// `usable_size` bytes at `ptr`. +/// +/// # Safety +/// +/// `ptr` is valid only during the call and must not be retained or deallocated. +/// The hook may allocate or free other allocations; nested allocator activity +/// is excluded from profiling. Hooks may run concurrently and must not unwind +/// across the `extern "C"` boundary. +pub type ProfSampleFreeHook = + unsafe extern "C" fn(ptr: *const c_void, usable_size: libc::size_t); + +/// Signature of a hook installable via [`set_prof_backtrace_hook`]. +/// +/// `jemalloc` invokes this hook to capture the stack trace for a sample. +/// +/// # Safety +/// +/// The pointers are valid only during the call and must not be retained. The +/// hook must write at most `max_length` frames into `backtrace`, store that +/// count through `backtrace_length`, support concurrent calls, and not unwind +/// across the `extern "C"` boundary. +pub type ProfBacktraceHook = unsafe extern "C" fn( + backtrace: *mut *mut c_void, + backtrace_length: *mut c_uint, + max_length: c_uint, +); + +/// Installs, replaces, or (with `None`) uninstalls the hook `jemalloc` +/// calls after deciding to sample an allocation, returning the +/// previously-installed hook. +/// +/// Corresponds to `experimental.hooks.prof_sample`. +/// +/// # Errors +/// +/// Returns an error (`ENOENT`) if `opt.prof` is `false` at runtime; see +/// [`prof_reset`]. Note that `opt.prof` being `true` is sufficient to +/// install a hook; [`prof_active`] need not be `true` (installing while +/// inactive is a no-op until activated). +/// +/// # Safety +/// +/// The caller must ensure the linked `jemalloc` uses the documented +/// [`ProfSampleHook`] ABI and that `hook`, if present, upholds its contract. +/// Replaced hooks may still be in flight, so their state must remain valid. +pub unsafe fn set_prof_sample_hook( + hook: Option, +) -> crate::error::Result> { + unsafe { crate::raw::update(b"experimental.hooks.prof_sample\0", hook) } +} + +/// Installs, replaces, or (with `None`) uninstalls the hook `jemalloc` +/// calls just before freeing a previously-sampled allocation, returning the +/// previously-installed hook. +/// +/// Corresponds to `experimental.hooks.prof_sample_free`. See +/// [`set_prof_sample_hook`] for the applicable error semantics. +/// +/// # Safety +/// +/// The caller must ensure the linked `jemalloc` uses the documented +/// [`ProfSampleFreeHook`] ABI and that `hook`, if present, upholds its contract. +/// Replaced hooks may still be in flight, so their state must remain valid. +pub unsafe fn set_prof_sample_free_hook( + hook: Option, +) -> crate::error::Result> { + unsafe { + crate::raw::update(b"experimental.hooks.prof_sample_free\0", hook) + } +} + +/// Installs or replaces the hook `jemalloc` calls to capture a sample's +/// backtrace, returning the previously-installed hook, if any. +/// +/// Corresponds to `experimental.hooks.prof_backtrace`. Unlike +/// [`set_prof_sample_hook`]/[`set_prof_sample_free_hook`], this hook cannot +/// be uninstalled (`jemalloc` rejects a `NULL` new hook with `EINVAL`). Install +/// [`noop_prof_backtrace_hook`] instead of `jemalloc`'s default unwinder if +/// backtraces aren't wanted. If present, the returned previous hook may be +/// restored later or invoked by the replacement during a valid backtrace-hook +/// call. +/// +/// # Errors +/// +/// Returns an error (`ENOENT`) if `opt.prof` is `false` at runtime; see +/// [`prof_reset`]. +/// +/// # Safety +/// +/// The caller must ensure the linked `jemalloc` uses the documented +/// [`ProfBacktraceHook`] ABI and that `hook` upholds its contract. Replaced +/// hooks may still be in flight, so their state must remain valid. +pub unsafe fn set_prof_backtrace_hook( + hook: ProfBacktraceHook, +) -> crate::error::Result> { + unsafe { + crate::raw::update(b"experimental.hooks.prof_backtrace\0", Some(hook)) + } +} + +/// A [`ProfBacktraceHook`] that reports an empty backtrace for every +/// sample. +/// +/// Installing this via [`set_prof_backtrace_hook`] disables `jemalloc`'s +/// own stack unwinding going forward: the per-allocation sampling +/// decision still happens at the configured rate (see [`lg_sample`]) +/// and [`set_prof_sample_hook`]/[`set_prof_sample_free_hook`] hooks still +/// fire, but with `backtrace_length` reported as `0`. Intended for +/// out-of-process samplers (e.g. an eBPF profiler) that capture their own +/// stacks and only need `jemalloc`'s sampling clock, since capturing a +/// backtrace it already unwinds itself is otherwise a per-sample cost +/// (page-aligned promotion, `tcache` bypass, and a `tdata` mutex are still +/// paid regardless of whether a backtrace is captured). +/// +/// # Safety +/// +/// Must only be invoked by `jemalloc` itself as an +/// `experimental.hooks.prof_backtrace` hook, which always passes a non-null +/// `backtrace_length`. +pub unsafe extern "C" fn noop_prof_backtrace_hook( + _backtrace: *mut *mut c_void, + backtrace_length: *mut c_uint, + _max_length: c_uint, +) { + *backtrace_length = 0; +} + +// Heap-allocates to force samples, so this needs a real allocator (`use_std`). +#[cfg(all(test, feature = "use_std"))] +mod hook_tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + static SAMPLE_HOOK_CALLS: AtomicUsize = AtomicUsize::new(0); + static SAMPLE_FREE_HOOK_CALLS: AtomicUsize = AtomicUsize::new(0); + static SAMPLE_BACKTRACE_LENGTH: AtomicUsize = AtomicUsize::new(usize::MAX); + + unsafe extern "C" fn counting_sample_hook( + _ptr: *const c_void, + _size: libc::size_t, + _backtrace: *mut *mut c_void, + backtrace_length: c_uint, + _usable_size: libc::size_t, + ) { + SAMPLE_BACKTRACE_LENGTH + .store(backtrace_length as usize, Ordering::SeqCst); + SAMPLE_HOOK_CALLS.fetch_add(1, Ordering::SeqCst); + } + + unsafe extern "C" fn counting_sample_free_hook( + _ptr: *const c_void, + _usable_size: libc::size_t, + ) { + SAMPLE_FREE_HOOK_CALLS.fetch_add(1, Ordering::SeqCst); + } + + // Exercises the whole hook contract end-to-end against real `jemalloc` + // ctls: activates profiling, resets the sampler to catch every + // allocation, installs counting hooks, allocates/frees, and asserts both + // hooks actually fired before restoring prior state. + #[test] + fn sample_and_sample_free_hooks_fire() { + // Requires `opt.prof`; see the `profiling` feature docs. The CI job + // for this test bakes `prof:true` via `JEMALLOC_SYS_WITH_MALLOC_CONF`. + if !prof::read().unwrap() { + return; + } + let was_active = prof_active::update(false).unwrap(); + let prev_lg_sample = lg_sample::read().unwrap(); + SAMPLE_BACKTRACE_LENGTH.store(usize::MAX, Ordering::SeqCst); + // lg_sample: 0 => average one sample per byte, i.e. every allocation. + // `jemalloc` only recomputes each thread's next sample distance + // (from the new `lg_sample`) once the current, already-primed + // distance (drawn under whatever `lg_sample` was in effect before + // this call, e.g. the crate's default of 512 KiB) has been + // exhausted, so the very next allocation isn't guaranteed to sample + // yet. Only allocations after that first one are. + prof_reset(0).unwrap(); + let prev_backtrace = + unsafe { set_prof_backtrace_hook(noop_prof_backtrace_hook) } + .unwrap() + .expect("jemalloc had no previous backtrace hook"); + let prev_sample = + unsafe { set_prof_sample_hook(Some(counting_sample_hook)) } + .unwrap(); + let prev_sample_free = unsafe { + set_prof_sample_free_hook(Some(counting_sample_free_hook)) + } + .unwrap(); + prof_active::write(true).unwrap(); + + // Warm up past any sample distance primed under the previous + // `lg_sample`. The default's geometric interval can reach ~18 MiB, so + // burn well past that before checking for samples. + for _ in 0..32 { + drop(Box::new([0u8; 1024 * 1024])); + } + + let before_sample = SAMPLE_HOOK_CALLS.load(Ordering::SeqCst); + let before_free = SAMPLE_FREE_HOOK_CALLS.load(Ordering::SeqCst); + for _ in 0..16 { + drop(Box::new([0u8; 4096])); + if SAMPLE_HOOK_CALLS.load(Ordering::SeqCst) > before_sample + && SAMPLE_FREE_HOOK_CALLS.load(Ordering::SeqCst) > before_free + { + break; + } + } + + prof_active::write(false).unwrap(); + unsafe { set_prof_sample_hook(prev_sample) }.unwrap(); + unsafe { set_prof_sample_free_hook(prev_sample_free) }.unwrap(); + unsafe { set_prof_backtrace_hook(prev_backtrace) }.unwrap(); + prof_reset(prev_lg_sample).unwrap(); + prof_active::write(was_active).unwrap(); + + assert!( + SAMPLE_HOOK_CALLS.load(Ordering::SeqCst) > before_sample, + "prof_sample hook did not fire after warm-up with lg_sample=0" + ); + assert!( + SAMPLE_FREE_HOOK_CALLS.load(Ordering::SeqCst) > before_free, + "prof_sample_free hook did not fire after freeing sampled allocations" + ); + assert_eq!(SAMPLE_BACKTRACE_LENGTH.load(Ordering::SeqCst), 0); + } +} diff --git a/jemalloc-sys/README.md b/jemalloc-sys/README.md index bce9ae000..e67a4e930 100644 --- a/jemalloc-sys/README.md +++ b/jemalloc-sys/README.md @@ -45,6 +45,24 @@ This crate provides following cargo feature flags: * `libgcc` (unless --disable-prof-libgcc) * `gcc intrinsics` (unless --disable-prof-gcc) + The matching `profiling` feature in `tikv-jemalloc-ctl` also exposes + `jemalloc`'s experimental `experimental.hooks.prof_sample`/ + `prof_sample_free`/`prof_backtrace` hooks, letting an external sampler (e.g. + an eBPF profiler) piggyback `jemalloc`'s sampling decision. To avoid + `jemalloc`'s own stack walk, install `noop_prof_backtrace_hook` through + `set_prof_backtrace_hook`. These hooks are not re-exported by + `tikv-jemallocator`. + + The feature compiles profiling support but does not enable profiling. To + enable it, configure `prof:true` before `jemalloc` initialises. This can be + done at process launch with the appropriate `MALLOC_CONF` environment + variable, typically `_RJEM_MALLOC_CONF` for prefixed builds. Use + `prof:true,prof_active:false` to install hooks before enabling sampling + through `prof.active`, or use `prof:true` to begin sampling immediately with + the default per-thread settings. Alternatively, + `JEMALLOC_SYS_WITH_MALLOC_CONF` can embed the same + configuration at build time. + * `profiling_libunwind` (configure `jemalloc` with `--enable-prof-libunwind`): Force jemalloc to use `libunwind` for backtracing during heap profiling instead of the default gcc-based unwinding, which has a