diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index a04e920db..656c05966 100644 --- a/bin/cli/src/main.rs +++ b/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}; diff --git a/prover/tests/calibration.rs b/prover/tests/calibration.rs index c7d4d66f5..a87f13175 100644 --- a/prover/tests/calibration.rs +++ b/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) diff --git a/prover/tests/jemalloc_conf.rs b/prover/tests/jemalloc_conf.rs new file mode 100644 index 000000000..b8661ce87 --- /dev/null +++ b/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::()`, 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" + ); +}