Add profiling_hooks feature exposing jemalloc's experimental sample hooks

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,
                     _ => (),
                 }
 

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<ProfSampleHook>,
+) -> crate::error::Result<Option<ProfSampleHook>> {
+    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<ProfSampleFreeHook>,
+) -> crate::error::Result<Option<ProfSampleFreeHook>> {
+    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<ProfBacktraceHook> {
+    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();
+    }
+}

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