perf(alloc): compile jemalloc's never-purge policy into the binary

bin/cli/src/main.rs

@@ -10,6 +10,55 @@ use clap::{Parser, Subcommand, ValueHint};
 
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+// jemalloc, never purging.
+//
+// The allocator itself is unchanged: jemalloc is here because the platform
+// allocator keeps freed arena chunks resident, and on the recursion campaign's
+// branch the same proves read up to 13 GiB higher under glibc. What this sets
+// is jemalloc's *decay* timers, which hand freed pages back to the OS. The
+// prover allocates and frees multi-hundred-MiB host buffers continuously, so
+// those pages come straight back as minor faults on the worker threads.
+//
+// Measured on an RTX 5090 box on the recursion campaign's branches, ABBA in
+// each, every arm at one commit and one set of knobs:
+//   * the WHIR prover (keccak, `whir/lfm` @ 64393da9) — 39.69-39.88 s a block
+//     with this setting against 43.94-44.07 s without, the default costing
+//     +13 M minor faults and +15 s of system time per run on this arm;
+//   * the per-table STARK tree (0e4f4610) — 187 / 163 / 166 / 171 s, both
+//     never-purge arms under both default arms, 5-24 s a block, with the proof
+//     bytes unmoved (30 identical lines, 0 differing).
+// The cost is peak RSS: +1.4 GiB and +2.9-3.3 GiB respectively, a ninth to a
+// quarter of the 13 GiB the allocator choice itself is worth — which is why the
+// lever is the decay setting and not the allocator. `background_thread:true` recovers
+// none of it: the cost is the re-touch, not the `madvise` call.
+//
+// This binary's own pipeline has not been measured under the setting; the
+// numbers above are from the campaign's branches, where the prover's
+// allocation pattern is the same.
+//
+// `_RJEM_MALLOC_CONF` in the environment still overrides this, which is how a
+// measurement arm puts the default policy back. It has to be that spelling:
+// `tikv-jemalloc-sys` builds with `--with-jemalloc-prefix=_rjem_` under default
+// features, and jemalloc then reads one env name chosen at configure time
+// (`jemalloc.c`, `obtain_malloc_conf` source 3) — so plain `MALLOC_CONF` is read
+// by nothing here and sets an arm to the default policy without saying it did
+// not. The file source is prefixed too: `/etc/_rjem_malloc.conf`.
+//
+// jemalloc reads this symbol as a `const char *` before `main` is entered, so
+// the value has to be in the initializer, and the name is the prefixed one
+// `tikv-jemalloc-sys` declares (`#[cfg_attr(prefixed, link_name =
+// "_rjem_malloc_conf")]`, its `src/lib.rs`). None of that is compiler-checked.
+// `prover/tests/jemalloc_conf.rs` reads both options back out of jemalloc, but
+// it carries its own copy of this block and reads its own process — it pins the
+// pattern, not this export. Deleting the lines below turns nothing red.
+const NEVER_PURGE: &[u8] = b"dirty_decay_ms:-1,muzzy_decay_ms:-1\0";
+
+#[allow(non_upper_case_globals)]
+#[unsafe(export_name = "_rjem_malloc_conf")]
+pub static malloc_conf: Option<&'static core::ffi::c_char> =
+    Some(unsafe { &*(NEVER_PURGE.as_ptr() as *const core::ffi::c_char) });
+
 use executor::vm::instruction::decoding::Instruction;
 use executor::vm::instruction::execution::{Accelerator, SyscallNumbers};
 use executor::{elf::Elf, flamegraph::FlamegraphGenerator, vm::execution::Executor};

prover/tests/calibration.rs

@@ -21,6 +21,22 @@ use tikv_jemalloc_ctl::{epoch, stats};
 #[global_allocator]
 static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 
+// ...with the shipped binary's purge policy, so this binary is the production
+// allocator *configuration* and not just the production allocator. The reason
+// and the numbers are at `bin/cli/src/main.rs`. `prover/tests/jemalloc_conf.rs`
+// asserts that this export pattern is read, but it does so against its own copy
+// in its own process — nothing checks the copy below.
+//
+// It moves nothing this file asserts — `stats::allocated` is live bytes, which
+// the decay timers do not touch; a resident-memory assertion added here later
+// would read the wrong configuration without it.
+const NEVER_PURGE: &[u8] = b"dirty_decay_ms:-1,muzzy_decay_ms:-1\0";
+
+#[allow(non_upper_case_globals)]
+#[unsafe(export_name = "_rjem_malloc_conf")]
+pub static malloc_conf: Option<&'static core::ffi::c_char> =
+    Some(unsafe { &*(NEVER_PURGE.as_ptr() as *const core::ffi::c_char) });
+
 fn allocated_bytes() -> usize {
     epoch::advance().ok();
     stats::allocated::read().unwrap_or(0)

prover/tests/jemalloc_conf.rs

@@ -0,0 +1,81 @@
+//! jemalloc's purge policy is compiled into the binary — checked by reading it
+//! back out of the allocator serving this process.
+//!
+//! `bin/cli/src/main.rs` and `calibration.rs` each export
+//! `_rjem_malloc_conf = "dirty_decay_ms:-1,muzzy_decay_ms:-1"` beside their
+//! `#[global_allocator]`, so the never-purge policy travels in the binary
+//! rather than in a launcher's environment. Nothing about that export is
+//! checked by the compiler: a misspelled symbol, a wrong value type, or a
+//! jemalloc built without the `_rjem_` prefix each leave a binary that
+//! compiles, links, runs — and quietly purges.
+//!
+//! This is its own test binary because the check needs a jemalloc process of
+//! its own: the prover's lib tests run under the platform allocator, where a
+//! `mallctl` read would say nothing, and `calibration.rs` is behind
+//! `disk-spill` and pays for a full proof. What it pins is the export pattern —
+//! symbol, type, initializer, edition spelling — in the copy below, which is
+//! byte-identical to the two production sites but not mechanically tied to
+//! them: delete either of those and this still passes. It is a self-test of the
+//! pattern, not a regression guard on the two sites that ship it. That the
+//! shipped `cli` binary carries the symbol is a link-time property, read with
+//! `nm` rather than asserted here.
+
+use tikv_jemalloc_ctl::raw;
+
+#[global_allocator]
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+/// The same string the two production sites export, character for character.
+const NEVER_PURGE: &[u8] = b"dirty_decay_ms:-1,muzzy_decay_ms:-1\0";
+
+#[allow(non_upper_case_globals)]
+#[unsafe(export_name = "_rjem_malloc_conf")]
+pub static malloc_conf: Option<&'static core::ffi::c_char> =
+    Some(unsafe { &*(NEVER_PURGE.as_ptr() as *const core::ffi::c_char) });
+
+/// jemalloc's default `opt.dirty_decay_ms`, quoted in the failure message so a
+/// red test says which value it found and where that value comes from.
+const DEFAULT_DIRTY_DECAY_MS: isize = 10_000;
+
+#[test]
+fn jemalloc_never_purge_is_compiled_in() {
+    // `_RJEM_MALLOC_CONF` sets these same options from the environment, and
+    // benchmark runs do set it; with it set, reading `-1` back would say nothing
+    // about the compiled-in export, so refuse to run rather than pass for the
+    // wrong reason. Plain `MALLOC_CONF` is inert in this prefixed build — guarded
+    // anyway, so that a future unprefixed build does not silently pass here.
+    // Not covered: `/etc/_rjem_malloc.conf`, the one remaining source that could
+    // set `opt.*` from outside this binary.
+    for var in ["_RJEM_MALLOC_CONF", "MALLOC_CONF"] {
+        assert!(
+            std::env::var_os(var).is_none(),
+            "{var} is set in this process's environment. jemalloc reads \
+             `_RJEM_MALLOC_CONF` (this build is prefixed), which sets `opt.*` on \
+             its own, so this test could not tell the compiled-in export from the \
+             environment; unset it and re-run."
+        );
+    }
+
+    // `opt.dirty_decay_ms` and `opt.muzzy_decay_ms` are jemalloc `ssize_t`s.
+    // `raw::read` asserts the mallctl's width equals `size_of::<T>()`, so a
+    // wrong Rust width fails here rather than reading a truncated value.
+    let dirty: isize =
+        unsafe { raw::read(b"opt.dirty_decay_ms\0") }.expect("opt.dirty_decay_ms is readable");
+    let muzzy: isize =
+        unsafe { raw::read(b"opt.muzzy_decay_ms\0") }.expect("opt.muzzy_decay_ms is readable");
+
+    assert_eq!(
+        dirty, -1,
+        "opt.dirty_decay_ms is {dirty}, not -1 (jemalloc's default is \
+         {DEFAULT_DIRTY_DECAY_MS}): the `_rjem_malloc_conf` export beside this \
+         file's `#[global_allocator]` is missing, misspelled, or was not read, \
+         and a binary built this way returns dirty pages to the OS on a timer"
+    );
+    assert_eq!(
+        muzzy, -1,
+        "opt.muzzy_decay_ms is {muzzy}, not -1 (jemalloc's default is 0): the \
+         `_rjem_malloc_conf` export beside this file's `#[global_allocator]` is \
+         missing, misspelled, or was not read, and a binary built this way \
+         unmaps muzzy pages immediately"
+    );
+}