From 46c352bcda651d494f7fc3ccc9e017f5c62d42cb Mon Sep 17 00:00:00 2001 From: Jaya Kasa Date: Tue, 16 Jun 2026 11:42:46 -0400 Subject: [PATCH] feat: Bazel integration cookbook + default_output_path helper Adds snmalloc-rs/docs/bazel.md cookbook covering the recommended profile-output path resolution chain (SNMALLOC_PROFILE_OUT > TEST_UNDECLARED_OUTPUTS_DIR > $TMPDIR/heap_{pid}.folded), BES upload size considerations, and an example rust_test snippet. Adds profile::default_output_path() in snmalloc-rs/src/profile.rs (gated on the profiling feature) that implements the chain at the bottom of the file -- outside the HeapProfile impl block so it merges cleanly past the in-flight write_flamegraph rename PR. README.md gains a single-line pointer to docs/bazel.md near the existing heap-profiling section. Verification: cargo build -p snmalloc-rs and cargo build -p snmalloc-rs --features profiling both clean; cargo test -p snmalloc-rs --features profiling --test profile_default_output_path covers the env-var precedence chain end to end. --- snmalloc-rs/README.md | 2 + snmalloc-rs/docs/bazel.md | 104 +++++++++++++++++ snmalloc-rs/src/profile.rs | 77 ++++++++++++ .../tests/profile_default_output_path.rs | 110 ++++++++++++++++++ 4 files changed, 293 insertions(+) create mode 100644 snmalloc-rs/docs/bazel.md create mode 100644 snmalloc-rs/tests/profile_default_output_path.rs diff --git a/snmalloc-rs/README.md b/snmalloc-rs/README.md index 876eac028..6dedc0089 100644 --- a/snmalloc-rs/README.md +++ b/snmalloc-rs/README.md @@ -41,6 +41,8 @@ There are the following features defined in this crate: ## Heap Profiling +See [`docs/bazel.md`](docs/bazel.md) for the Bazel-integration cookbook (profile-output path resolution, BES upload limits, opt-in `rust_test` snippet). + The `profiling` Cargo feature enables a low-overhead statistical heap profiler in the underlying snmalloc build. Each allocation has an independent Poisson probability of being recorded with its call stack; diff --git a/snmalloc-rs/docs/bazel.md b/snmalloc-rs/docs/bazel.md new file mode 100644 index 000000000..f165e5a8b --- /dev/null +++ b/snmalloc-rs/docs/bazel.md @@ -0,0 +1,104 @@ +# Bazel integration cookbook + +This page collects the patterns we use to run snmalloc-rs heap-profile +collection from inside a Bazel-built binary or `rust_test`. It assumes +familiarity with the `profiling` Cargo feature (see the main +`snmalloc-rs/README.md` for the API surface). + +## Profile-output path resolution + +`snmalloc_rs::profile::default_output_path()` (gated on the `profiling` +feature) returns a `PathBuf` chosen by the following precedence chain. +First match wins: + +1. **`SNMALLOC_PROFILE_OUT`** — explicit override. Whatever the + operator / CI script puts here is used verbatim. This is the escape + hatch you want to wire into a `--test_env=` flag (Bazel test) or a + `--action_env=` flag (Bazel binary) so you can redirect output + without recompiling the workload. +2. **`$TEST_UNDECLARED_OUTPUTS_DIR/heap.folded`** — Bazel's + per-test scratch directory. When a `rust_test` runs under + `bazel test`, Bazel sets `TEST_UNDECLARED_OUTPUTS_DIR` to a + per-action directory and automatically picks up anything written + there as a declared test artefact. The file ends up in the + `outputs.zip` attached to the test result and is visible in the + Build Event Service (BES) / Remote Build Execution (RBE) UI without + any extra wiring. +3. **`$TMPDIR/heap_{pid}.folded`** — final fallback for plain + `cargo run` / `cargo test` / interactive `bazel run`. The PID is + appended so two concurrent processes don't clobber each other's + output. + +The three rungs are intentionally ordered from "most explicit" to +"safest default". The `SNMALLOC_PROFILE_OUT` override is the only one +that respects the literal path you set — both of the other rungs +synthesize a filename for you. Callers writing pprof or another +format should `with_extension(...)` the returned path; the default +suffix is `.folded` because that's the most broadly consumable format +emitted by `HeapProfile::write_flamegraph`. + +## BES upload size considerations + +Bazel's Build Event Service uploads test artefacts back to the +result-store on every `bazel test` invocation. The default per-file +upload size cap is around **10 MiB**; profiles larger than that will +either be truncated or rejected depending on the BES backend. A few +practical implications: + +- A 512 KiB sampling rate (the snmalloc default) keeps a folded-stack + profile well under the cap for workloads up to a few minutes of + steady-state allocation. If your workload runs longer, raise the + sampling rate (set `SNMALLOC_PROFILE_RATE=2097152` for 2 MiB, etc.) + to keep the output bounded. +- For very long-running test workloads, rotate the output: take a + snapshot every N seconds, write it to a numbered file under + `$TEST_UNDECLARED_OUTPUTS_DIR/heap_{N}.folded`, and let downstream + tooling stitch them. The `default_output_path` helper only resolves + a single path; rotation is a one-line `with_file_name()` away. +- Gzipped pprof (`HeapProfile::write_pprof_gz`) typically shrinks + output 5–10x versus the folded form. If you're already collecting + pprof, prefer the gzipped variant for the BES round-trip. + +## Example `BUILD.bazel` snippet + +The minimal opt-in pattern for a `rust_test` that wants to dump a +heap profile on exit: + +```python +load("@rules_rust//rust:defs.bzl", "rust_test") + +rust_test( + name = "my_heap_profile_test", + srcs = ["tests/my_heap_profile_test.rs"], + edition = "2021", + deps = [ + "//snmalloc-rs:snmalloc_rs", + ], + # Opt the test into snmalloc's heap profiler at a 256 KiB + # sampling rate. `SNMALLOC_PROFILE_OUT` is left unset so the + # path-resolution chain falls through to TEST_UNDECLARED_OUTPUTS_DIR, + # which Bazel auto-uploads as a test artefact. + env = { + "SNMALLOC_PROFILE_ENABLE": "1", + "SNMALLOC_PROFILE_RATE": "262144", + }, +) +``` + +If you want to override the path explicitly — e.g. to dump to a known +location for a downstream `genrule` to consume — extend `env` with +`SNMALLOC_PROFILE_OUT`: + +```python + env = { + "SNMALLOC_PROFILE_ENABLE": "1", + "SNMALLOC_PROFILE_RATE": "262144", + "SNMALLOC_PROFILE_OUT": "/tmp/explicit_heap.folded", + }, +``` + +For a `rust_binary` invoked via `bazel run`, swap `env` for the same +keys on a wrapper `sh_binary` or pass `--action_env=...` on the +command line. The resolution chain in `default_output_path()` is +identical regardless of the host rule kind — the helper only inspects +the process environment at call time. diff --git a/snmalloc-rs/src/profile.rs b/snmalloc-rs/src/profile.rs index a212674dd..0cd13256d 100644 --- a/snmalloc-rs/src/profile.rs +++ b/snmalloc-rs/src/profile.rs @@ -1542,6 +1542,83 @@ impl SnMalloc { } } +/// Resolve a default filesystem path to write a serialised heap +/// profile (folded-stack, pprof, etc.) to, using the precedence chain +/// recommended for Bazel + dev integrations. See +/// `snmalloc-rs/docs/bazel.md` for the cookbook explaining the +/// rationale for each fallback step. +/// +/// Precedence (first match wins): +/// +/// 1. `SNMALLOC_PROFILE_OUT` -- explicit override. Always honoured +/// verbatim; lets operators / CI scripts redirect output without +/// recompiling. +/// 2. `TEST_UNDECLARED_OUTPUTS_DIR` -- Bazel's per-test scratch +/// directory. When set, the file is written as +/// `$TEST_UNDECLARED_OUTPUTS_DIR/heap.folded` so that Bazel +/// automatically uploads it as a declared test output (visible in +/// BES / RBE result UIs). +/// 3. `std::env::temp_dir()` -- final fallback for plain `cargo run` +/// / `cargo test` invocations. The PID is appended +/// (`heap_{pid}.folded`) so concurrent processes don't clobber each +/// other. +/// +/// The returned path is intentionally `.folded`-suffixed -- this is +/// the most broadly consumable format produced by +/// [`HeapProfile::write_flamegraph`] / [`HeapProfile::write_flamegraph_with`]. +/// Callers writing pprof or another format should `with_extension` +/// the returned path. +/// +/// Only available with the `profiling` Cargo feature. +/// +/// # Example +/// +/// ```no_run +/// # #[cfg(feature = "profiling")] +/// # fn main() -> std::io::Result<()> { +/// use snmalloc_rs::SnMalloc; +/// use snmalloc_rs::profile::default_output_path; +/// use std::fs::File; +/// +/// let profile = SnMalloc.snapshot(); +/// let path = default_output_path(); +/// let mut f = File::create(&path)?; +/// profile.write_flamegraph(&mut f)?; +/// # Ok(()) +/// # } +/// # #[cfg(not(feature = "profiling"))] +/// # fn main() {} +/// ``` +#[cfg(feature = "profiling")] +pub fn default_output_path() -> std::path::PathBuf { + // 1. Explicit override wins. An empty string is treated as + // "unset" so a stray `SNMALLOC_PROFILE_OUT=` in a shell + // profile doesn't accidentally point us at the current + // directory. + if let Ok(p) = std::env::var("SNMALLOC_PROFILE_OUT") { + if !p.is_empty() { + return std::path::PathBuf::from(p); + } + } + // 2. Bazel sets TEST_UNDECLARED_OUTPUTS_DIR per + // https://bazel.build/reference/test-encyclopedia#initial-conditions + // so any file written there is uploaded by Bazel as a + // declared test artefact. + if let Ok(dir) = std::env::var("TEST_UNDECLARED_OUTPUTS_DIR") { + if !dir.is_empty() { + let mut p = std::path::PathBuf::from(dir); + p.push("heap.folded"); + return p; + } + } + // 3. Final fallback for plain `cargo run` / `cargo test` / + // interactive use. Stamp the PID so concurrent runs don't + // overwrite each other. + let mut p = std::env::temp_dir(); + p.push(std::format!("heap_{}.folded", std::process::id())); + p +} + #[cfg(test)] mod tests { use super::*; diff --git a/snmalloc-rs/tests/profile_default_output_path.rs b/snmalloc-rs/tests/profile_default_output_path.rs new file mode 100644 index 000000000..f9103d61a --- /dev/null +++ b/snmalloc-rs/tests/profile_default_output_path.rs @@ -0,0 +1,110 @@ +//! Tests for `snmalloc_rs::profile::default_output_path` -- the +//! Bazel-aware path-resolution helper introduced in ticket +//! 86aj2dwrr. The helper inspects three process-global environment +//! variables; we exercise the precedence chain end-to-end here. +//! +//! These tests are gated on the `profiling` Cargo feature because the +//! helper itself only exists in that build configuration. Without +//! the feature this file compiles down to an empty `tests` binary. +//! +//! All env-var manipulation runs in a single `#[test]` so the +//! save/restore is locally serialised; spawning multiple tests would +//! race against each other on the shared environment. + +#![cfg(feature = "profiling")] + +use snmalloc_rs::profile::default_output_path; +use std::env; +use std::path::PathBuf; + +const ENV_OUT: &str = "SNMALLOC_PROFILE_OUT"; +const ENV_BAZEL: &str = "TEST_UNDECLARED_OUTPUTS_DIR"; + +/// Save the current value of an env var, return a guard that restores +/// it on drop. This keeps the test idempotent w.r.t. the surrounding +/// process environment -- important because `cargo test` runs all +/// integration tests in a single binary and may set +/// `TEST_UNDECLARED_OUTPUTS_DIR` itself in some CI configurations. +struct EnvGuard { + key: &'static str, + prior: Option, +} + +impl EnvGuard { + fn save(key: &'static str) -> Self { + let prior = env::var(key).ok(); + Self { key, prior } + } +} + +impl Drop for EnvGuard { + fn drop(&mut self) { + match &self.prior { + Some(v) => env::set_var(self.key, v), + None => env::remove_var(self.key), + } + } +} + +#[test] +fn precedence_chain_exhaustive() { + // Save originals so concurrent test binaries / parent env are + // restored on exit. EnvGuard restores on drop in reverse + // declaration order. + let _g_out = EnvGuard::save(ENV_OUT); + let _g_bazel = EnvGuard::save(ENV_BAZEL); + + // -------- 1. Explicit override wins ------------------------------- + env::set_var(ENV_OUT, "/tmp/explicit.folded"); + env::set_var(ENV_BAZEL, "/tmp/bazel_should_be_ignored"); + let p = default_output_path(); + assert_eq!( + p, + PathBuf::from("/tmp/explicit.folded"), + "SNMALLOC_PROFILE_OUT must take precedence verbatim" + ); + + // Empty SNMALLOC_PROFILE_OUT is treated as unset so a stray + // `SNMALLOC_PROFILE_OUT=` in a shell profile doesn't pin us to + // the current working directory. + env::set_var(ENV_OUT, ""); + env::set_var(ENV_BAZEL, "/tmp/bazel_outputs"); + let p = default_output_path(); + assert_eq!( + p, + PathBuf::from("/tmp/bazel_outputs/heap.folded"), + "empty SNMALLOC_PROFILE_OUT must fall through to Bazel path" + ); + + // -------- 2. Bazel TEST_UNDECLARED_OUTPUTS_DIR rung ---------------- + env::remove_var(ENV_OUT); + env::set_var(ENV_BAZEL, "/tmp/bazel_outputs"); + let p = default_output_path(); + assert_eq!( + p, + PathBuf::from("/tmp/bazel_outputs/heap.folded"), + "TEST_UNDECLARED_OUTPUTS_DIR must be suffixed with heap.folded" + ); + + // -------- 3. tmp_dir / pid fallback -------------------------------- + env::remove_var(ENV_OUT); + env::remove_var(ENV_BAZEL); + let p = default_output_path(); + // Final rung lives under env::temp_dir() and the file name carries + // the current PID. Both invariants matter -- the temp-dir prefix + // ensures we never accidentally write into the source tree, and + // the PID stamp prevents concurrent processes from racing on the + // same path. + let tmp = env::temp_dir(); + assert!( + p.starts_with(&tmp), + "fallback path {p:?} must live under temp_dir {tmp:?}" + ); + let fname = p + .file_name() + .expect("fallback path has a file name") + .to_str() + .expect("file name is valid utf-8"); + let expected = format!("heap_{}.folded", std::process::id()); + assert_eq!(fname, expected, "fallback file name must encode the PID"); +}