feat: Bazel integration cookbook + default_output_path helper

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;

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.

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::*;

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<String>,
+}
+
+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");
+}