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..1c878334c 100644 --- a/jemalloc-ctl/src/profiling.rs +++ b/jemalloc-ctl/src/profiling.rs @@ -1,6 +1,12 @@ //! `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`]. option! { lg_prof_interval[ str: b"opt.lg_prof_interval\0", non_str: 2 ] => libc::ssize_t | @@ -153,3 +159,292 @@ 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`. + /// + /// Note: `opt.prof_thread_active_init` is unrelated. It controls the + /// per-thread `thread.prof.active` flag, not this global toggle. + /// + /// While inactive, sampling hooks installed via the `profiling` + /// 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. 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(lg_sample: libc::size_t) -> crate::error::Result<()> { + unsafe { crate::raw::write(b"prof.reset\0", lg_sample) } +} + +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 an array of `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. It may itself +/// allocate or free; nested allocator activity is excluded from profiling. +/// It 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`. See [`ProfSampleHook`] for the applicable safety +/// contract. +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. +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). +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. +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. 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`]. +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`. +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}; + + union Conf { + bytes: &'static u8, + c_char: &'static libc::c_char, + } + + // Enable profiling only for this test binary. Normal library builds leave + // the process-wide profiling policy to the final consumer. + #[export_name = "_rjem_malloc_conf"] + pub static TEST_MALLOC_CONF: Option<&'static libc::c_char> = + Some(unsafe { + Conf { + bytes: &b"prof:true,prof_active:false\0"[0], + } + .c_char + }); + + 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/README.md b/jemalloc-sys/README.md index bce9ae000..f4e6ad075 100644 --- a/jemalloc-sys/README.md +++ b/jemalloc-sys/README.md @@ -45,6 +45,21 @@ 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 without its stack + walking. 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. + 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