From ff0e303b5b03f4215511b805916fc9abcab9292f Mon Sep 17 00:00:00 2001 From: Scott Gerring Date: Sat, 11 Jul 2026 16:56:21 +0200 Subject: [PATCH] Add profiling_hooks feature exposing jemalloc's experimental sample hooks Signed-off-by: Scott Gerring --- jemalloc-ctl/Cargo.toml | 3 +- jemalloc-ctl/src/macros.rs | 2 + jemalloc-ctl/src/profiling.rs | 291 ++++++++++++++++++++++++++++++++++ jemalloc-sys/Cargo.toml | 1 + jemalloc-sys/README.md | 14 ++ jemalloc-sys/build.rs | 17 ++ 6 files changed, 327 insertions(+), 1 deletion(-) diff --git a/jemalloc-ctl/Cargo.toml b/jemalloc-ctl/Cargo.toml index 1f93c84a4..deea2a2fa 100644 --- a/jemalloc-ctl/Cargo.toml +++ b/jemalloc-ctl/Cargo.toml @@ -38,9 +38,10 @@ default = [] stats = ["tikv-jemalloc-sys/stats"] profiling = ["tikv-jemalloc-sys/profiling"] profiling_libunwind = ["tikv-jemalloc-sys/profiling_libunwind", "profiling"] +profiling_hooks = ["tikv-jemalloc-sys/profiling_hooks", "profiling"] use_std = [ "libc/use_std" ] disable_initial_exec_tls = ["tikv-jemalloc-sys/disable_initial_exec_tls"] [package.metadata.docs.rs] rustdoc-args = [ "--cfg", "jemallocator_docs" ] -features = ["stats", "profiling", "use_std"] +features = ["stats", "profiling", "profiling_hooks", "use_std"] diff --git a/jemalloc-ctl/src/macros.rs b/jemalloc-ctl/src/macros.rs index 94050e354..1eb5b777a 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, + // Skipped: races with the hook test's own `prof_active` writes. + "prof_active" if cfg!(feature = "profiling_hooks") => return, _ => (), } diff --git a/jemalloc-ctl/src/profiling.rs b/jemalloc-ctl/src/profiling.rs index 306ab1a8e..28778a036 100644 --- a/jemalloc-ctl/src/profiling.rs +++ b/jemalloc-ctl/src/profiling.rs @@ -1,6 +1,17 @@ //! `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`]. +#![cfg_attr( + feature = "profiling_hooks", + doc = " +With the `profiling_hooks` feature, it additionally exposes `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`]." +)] option! { lg_prof_interval[ str: b"opt.lg_prof_interval\0", non_str: 2 ] => libc::ssize_t | @@ -153,3 +164,283 @@ 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`] always returns `false` and writing to it fails with + /// `ENOENT`. `jemalloc` initializes [`prof_active`] to + /// `opt.prof_thread_active_init` (which itself defaults to `true`) as + /// soon as `opt.prof` is `true`, so a build with `opt.prof` enabled + /// samples by default unless [`prof_active`] is set to `false`, e.g. via + /// `prof_active:false` in `MALLOC_CONF`. + /// + /// While inactive, sampling hooks installed via the `profiling_hooks` + /// feature's hook setters remain installed but do not fire, since no + /// allocation is ever selected for sampling. + /// + /// # 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::write(false).unwrap(); + /// # let _ = was_active; + /// # } + /// ``` + mib_docs: /// See [`prof_active`]. +} + +/// Resets `jemalloc`'s heap profile sample accumulators and, going forward, +/// samples allocations at a rate of one per `2^lg_sample` bytes of +/// allocation activity. +/// +/// 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, e.g. +/// because `MALLOC_CONF`/`JEMALLOC_SYS_WITH_MALLOC_CONF` overrode the +/// `prof:true` this crate's `profiling`/`profiling_hooks` features bake in. +pub fn prof_reset(lg_sample: libc::size_t) -> crate::error::Result<()> { + unsafe { crate::raw::write(b"prof.reset\0", lg_sample) } +} + +#[cfg(feature = "profiling_hooks")] +use libc::{c_uint, c_void}; + +/// 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 `backtrace_length` +/// `void*` frames captured by the installed [`ProfBacktraceHook`] +/// (`backtrace_length` is `0` if [`noop_prof_backtrace_hook`] is installed). +/// +/// # Safety +/// +/// No `jemalloc` mutex is held while this hook runs, but it does run inside +/// `jemalloc`'s `pre_reentrancy`/`post_reentrancy` bracket and may itself +/// allocate/free (including recursively sampling). It must not unwind +/// across the `extern "C"` boundary — a panic here is undefined behavior +/// below Rust 1.81, and this crate is `no_std` by default so `catch_unwind` +/// is unavailable regardless. +#[cfg(feature = "profiling_hooks")] +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`. See [`ProfSampleHook`] for the applicable safety +/// contract. +#[cfg(feature = "profiling_hooks")] +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; it +/// must write at most `max_length` frames into `backtrace` and store the +/// number of frames written through `backtrace_length`. See +/// [`ProfSampleHook`] for the applicable safety contract. +#[cfg(feature = "profiling_hooks")] +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). +#[cfg(feature = "profiling_hooks")] +pub 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. +#[cfg(feature = "profiling_hooks")] +pub 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. +/// +/// 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, and keep the returned previous +/// hook only to log/inspect, not to call directly, since it may itself be a +/// previously-installed [`noop_prof_backtrace_hook`] or `jemalloc`'s +/// internal default depending on prior state. +/// +/// # Errors +/// +/// Returns an error (`ENOENT`) if `opt.prof` is `false` at runtime — see +/// [`prof_reset`]. +#[cfg(feature = "profiling_hooks")] +pub fn set_prof_backtrace_hook( + hook: ProfBacktraceHook, +) -> crate::error::Result { + unsafe { crate::raw::update(b"experimental.hooks.prof_backtrace\0", 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_prof_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`. +#[cfg(feature = "profiling_hooks")] +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 = "profiling_hooks", 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); + + 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_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() { + let was_active = prof_active::read().unwrap(); + let prev_lg_sample = lg_prof_sample::read().unwrap(); + // 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(); + prof_active::write(true).unwrap(); + + let prev_sample = + set_prof_sample_hook(Some(counting_sample_hook)).unwrap(); + let prev_sample_free = + set_prof_sample_free_hook(Some(counting_sample_free_hook)) + .unwrap(); + + // Warm up past any stale pre-reset sample distance: 16 MiB is many + // times the largest plausible leftover distance from a 512 KiB mean. + for _ in 0..16 { + 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; + } + } + + set_prof_sample_hook(prev_sample).unwrap(); + set_prof_sample_free_hook(prev_sample_free).unwrap(); + prof_active::write(was_active).unwrap(); + prof_reset(prev_lg_sample).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" + ); + } + + #[test] + fn backtrace_hook_can_be_replaced_and_restored() { + let prev = set_prof_backtrace_hook(noop_prof_backtrace_hook).unwrap(); + set_prof_backtrace_hook(prev).unwrap(); + } +} diff --git a/jemalloc-sys/Cargo.toml b/jemalloc-sys/Cargo.toml index 9b3d922f1..53d3314b6 100644 --- a/jemalloc-sys/Cargo.toml +++ b/jemalloc-sys/Cargo.toml @@ -42,6 +42,7 @@ cc = "^1.0.13" default = ["background_threads_runtime_support"] profiling = [] profiling_libunwind = ["profiling"] +profiling_hooks = ["profiling"] debug = [] background_threads_runtime_support = [] background_threads = [ "background_threads_runtime_support" ] diff --git a/jemalloc-sys/README.md b/jemalloc-sys/README.md index bce9ae000..c590a1b68 100644 --- a/jemalloc-sys/README.md +++ b/jemalloc-sys/README.md @@ -54,6 +54,20 @@ This crate provides following cargo feature flags: to be installed. On macOS/iOS, unwind symbols are provided by the system and no extra library is needed. +* `profiling_hooks`: Enables `profiling` automatically and bakes + `prof:true,prof_active:false` into the default `malloc_conf`, so sampling + is installable but inert until a consumer flips `prof.active` at runtime. + Exposes `jemalloc`'s experimental + `experimental.hooks.prof_sample`/`prof_sample_free`/`prof_backtrace` hooks + through `tikv-jemalloc-ctl`'s `profiling` module, letting an external + sampler (e.g. an eBPF profiler) piggyback `jemalloc`'s sampling decision + without its stack walking. Not re-exported by `tikv-jemallocator`. + + Since this crate has `links = "jemalloc"`, there's only one `jemalloc` + build per dependency graph, so enabling this on any dependent applies + `prof:true,prof_active:false` to that shared build for everyone in the + graph too. + * `stats` (configure `jemalloc` with `--enable-stats`): Enable statistics gathering functionality. See the `jemalloc`'s "`opt.stats_print`" option documentation for usage details. diff --git a/jemalloc-sys/build.rs b/jemalloc-sys/build.rs index 845b9d43f..a2a797c85 100644 --- a/jemalloc-sys/build.rs +++ b/jemalloc-sys/build.rs @@ -246,6 +246,23 @@ fn main() { malloc_conf += "background_thread:false"; } + if env::var("CARGO_FEATURE_PROFILING_HOOKS").is_ok() { + info!("CARGO_FEATURE_PROFILING_HOOKS set"); + // Baking bare `prof:true` would also default `prof_active` and + // `prof_thread_active_init` to true. Since `links = "jemalloc"` + // unifies this build across every dependent in the graph, that would + // make the one shared jemalloc actively sample every allocation + // process-wide as soon as any dependent enables this feature, even + // if no hook is ever installed. Keep sampling installable + // (`opt_prof`) but inert until a consumer flips `prof.active` (or + // calls `tikv_jemalloc_ctl::profiling::prof_active::write(true)`) at + // runtime. + if !malloc_conf.is_empty() { + malloc_conf.push(','); + } + malloc_conf.push_str("prof:true,prof_active:false"); + } + if let Ok(malloc_conf_opts) = read_and_watch_env("JEMALLOC_SYS_WITH_MALLOC_CONF") { if !malloc_conf.is_empty() { malloc_conf.push(',');