Add profiling_hooks feature exposing jemalloc's experimental sample hooks

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"]

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

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<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.
+#[cfg(feature = "profiling_hooks")]
+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, 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<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`.
+#[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();
+    }
+}

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" ]

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.

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(',');