From d82cd273a630559de98f9adfa7adb7425e4e5356 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 14:31:41 -0700 Subject: [PATCH 1/6] Add jemalloc heap profiling support to moq-relay Optional `jemalloc` cargo feature that replaces the system allocator with jemalloc and enables heap profiling. Send SIGUSR1 to dump a profile to /tmp/moq-relay.heap. for analysis with jeprof. Enabled by default in Nix builds with MALLOC_CONF set in the systemd unit. Co-Authored-By: Claude Opus 4.6 (1M context) --- Cargo.lock | 33 ++++++++++++++++++++++++ nix/modules/moq-relay.nix | 3 +++ nix/overlay.nix | 2 +- rs/moq-relay/Cargo.toml | 3 +++ rs/moq-relay/src/jemalloc_prof.rs | 42 +++++++++++++++++++++++++++++++ rs/moq-relay/src/lib.rs | 2 ++ rs/moq-relay/src/main.rs | 7 ++++++ 7 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 rs/moq-relay/src/jemalloc_prof.rs diff --git a/Cargo.lock b/Cargo.lock index 7eea54a3ce..40eff546a4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3587,6 +3587,8 @@ dependencies = [ "serde_with", "tempfile", "thiserror 2.0.18", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "tokio", "toml", "tower-http", @@ -6420,6 +6422,37 @@ dependencies = [ "trackable", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "661f1f6a57b3a36dc9174a2c10f19513b4866816e13425d3e418b11cc37bc24c" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + +[[package]] +name = "tikv-jemalloc-sys" +version = "0.6.1+5.3.0-1-ge13ca993e8ccb9ba9847cc330696e02839f328f7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd8aa5b2ab86a2cefa406d889139c162cbb230092f7d1d7cbc1716405d852a3b" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0359b4327f954e0567e69fb191cf1436617748813819c94b8cd4a431422d053a" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.47" diff --git a/nix/modules/moq-relay.nix b/nix/modules/moq-relay.nix index 2e3b7287d9..20b1d2e8b4 100644 --- a/nix/modules/moq-relay.nix +++ b/nix/modules/moq-relay.nix @@ -207,6 +207,9 @@ in }; environment = { + # Enable jemalloc heap profiling; dump with `kill -USR1 ` + MALLOC_CONF = "prof:true,prof_active:true,prof_prefix:/tmp/moq-relay.heap"; + MOQ_LOG_LEVEL = lib.mkDefault cfg.logLevel; # Server configuration diff --git a/nix/overlay.nix b/nix/overlay.nix index 5bd2ab0161..abe7d32cc9 100644 --- a/nix/overlay.nix +++ b/nix/overlay.nix @@ -12,7 +12,7 @@ in crateInfo ../rs/moq-relay/Cargo.toml // { src = craneLib.cleanCargoSource ../.; - cargoExtraArgs = "-p moq-relay"; + cargoExtraArgs = "-p moq-relay --features jemalloc"; # Enable frame pointers for profiling support (negligible overhead on x86_64). # This also ensures the CDN build matches what Cachix caches. RUSTFLAGS = "-C force-frame-pointers=yes"; diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index 1e0152ccf7..75a14951bb 100644 --- a/rs/moq-relay/Cargo.toml +++ b/rs/moq-relay/Cargo.toml @@ -23,6 +23,7 @@ doc = false [features] default = ["iroh", "quinn", "websocket"] iroh = ["moq-native/iroh"] +jemalloc = ["dep:tikv-jemallocator", "dep:tikv-jemalloc-ctl"] noq = ["moq-native/noq"] quinn = ["moq-native/quinn"] quiche = ["moq-native/quiche"] @@ -49,6 +50,8 @@ thiserror = "2" tokio = { workspace = true, features = ["full"] } toml = "0.9" tower-http = { version = "0.6", features = ["cors"] } +tikv-jemallocator = { version = "0.6", features = ["profiling", "unprefixed_malloc_on_supported_platforms"], optional = true } +tikv-jemalloc-ctl = { version = "0.6", optional = true } tracing = "0.1" url = { version = "2", features = ["serde"] } diff --git a/rs/moq-relay/src/jemalloc_prof.rs b/rs/moq-relay/src/jemalloc_prof.rs new file mode 100644 index 0000000000..65ac812c58 --- /dev/null +++ b/rs/moq-relay/src/jemalloc_prof.rs @@ -0,0 +1,42 @@ +use tikv_jemalloc_ctl::raw; + +/// Activate jemalloc heap profiling and spawn a SIGUSR1 signal handler to dump profiles. +/// +/// Send `kill -USR1 ` to write a heap profile to `/tmp/moq-relay.heap.`. +pub fn init() { + let prof_active = b"prof.active\0"; + + match unsafe { raw::read::(prof_active) } { + Ok(true) => tracing::info!("jemalloc heap profiling is active"), + Ok(false) => { + tracing::info!("jemalloc profiling compiled in; activating"); + unsafe { raw::write(prof_active, true) }.ok(); + } + Err(err) => { + tracing::warn!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable"); + return; + } + } + + tokio::spawn(async { + if let Err(err) = signal_handler().await { + tracing::error!(%err, "jemalloc signal handler failed"); + } + }); +} + +async fn signal_handler() -> anyhow::Result<()> { + let mut sig = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::user_defined1())?; + let mut seq = 0u64; + + loop { + sig.recv().await; + let path = format!("/tmp/moq-relay.heap.{seq}\0"); + seq += 1; + + match unsafe { raw::write(b"prof.dump\0", path.as_ptr()) } { + Ok(()) => tracing::info!(path = &path[..path.len() - 1], "heap profile dumped"), + Err(err) => tracing::error!(%err, "failed to dump heap profile"), + } + } +} diff --git a/rs/moq-relay/src/lib.rs b/rs/moq-relay/src/lib.rs index 455c4fc4cc..b8fae4285b 100644 --- a/rs/moq-relay/src/lib.rs +++ b/rs/moq-relay/src/lib.rs @@ -11,6 +11,8 @@ mod auth; mod cluster; mod config; mod connection; +#[cfg(feature = "jemalloc")] +pub mod jemalloc_prof; mod web; #[cfg(feature = "websocket")] mod websocket; diff --git a/rs/moq-relay/src/main.rs b/rs/moq-relay/src/main.rs index 4bfba4754b..14d36d0f62 100644 --- a/rs/moq-relay/src/main.rs +++ b/rs/moq-relay/src/main.rs @@ -2,8 +2,15 @@ use moq_relay::*; use anyhow::Context; +#[cfg(feature = "jemalloc")] +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + #[tokio::main] async fn main() -> anyhow::Result<()> { + #[cfg(feature = "jemalloc")] + jemalloc_prof::init(); + // TODO: It would be nice to remove this and rely on feature flags only. // However, some dependency is pulling in `ring` and I don't know why, so meh for now. rustls::crypto::aws_lc_rs::default_provider() From cd778031e2054dbd3c233b749f461e75b9f76ae2 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 15:09:36 -0700 Subject: [PATCH 2/6] Refactor jemalloc profiling: async run pattern, configurable Nix dump path - Replace spawned task with async run() in tokio::select! like other modules - Delegate dump path to jemalloc's prof_prefix via MALLOC_CONF - Add configurable heapDumpPrefix option to the Nix module - Rename jemalloc_prof.rs to jemalloc.rs Co-Authored-By: Claude Opus 4.6 (1M context) --- nix/modules/moq-relay.nix | 8 +++++- .../src/{jemalloc_prof.rs => jemalloc.rs} | 25 ++++++------------- rs/moq-relay/src/lib.rs | 2 +- rs/moq-relay/src/main.rs | 9 ++++--- 4 files changed, 22 insertions(+), 22 deletions(-) rename rs/moq-relay/src/{jemalloc_prof.rs => jemalloc.rs} (51%) diff --git a/nix/modules/moq-relay.nix b/nix/modules/moq-relay.nix index 20b1d2e8b4..4bbbc71a91 100644 --- a/nix/modules/moq-relay.nix +++ b/nix/modules/moq-relay.nix @@ -142,6 +142,12 @@ in default = "/var/lib/moq-relay"; description = "State directory for keys and runtime data"; }; + + heapDumpPrefix = lib.mkOption { + type = lib.types.str; + default = "/tmp/moq-relay.heap"; + description = "Path prefix for jemalloc heap profile dumps (triggered via kill -USR1)"; + }; }; config = lib.mkIf cfg.enable { @@ -208,7 +214,7 @@ in environment = { # Enable jemalloc heap profiling; dump with `kill -USR1 ` - MALLOC_CONF = "prof:true,prof_active:true,prof_prefix:/tmp/moq-relay.heap"; + MALLOC_CONF = "prof:true,prof_active:true,prof_prefix:${cfg.heapDumpPrefix}"; MOQ_LOG_LEVEL = lib.mkDefault cfg.logLevel; diff --git a/rs/moq-relay/src/jemalloc_prof.rs b/rs/moq-relay/src/jemalloc.rs similarity index 51% rename from rs/moq-relay/src/jemalloc_prof.rs rename to rs/moq-relay/src/jemalloc.rs index 65ac812c58..fdc89c81d2 100644 --- a/rs/moq-relay/src/jemalloc_prof.rs +++ b/rs/moq-relay/src/jemalloc.rs @@ -1,9 +1,10 @@ use tikv_jemalloc_ctl::raw; -/// Activate jemalloc heap profiling and spawn a SIGUSR1 signal handler to dump profiles. +/// Activate jemalloc heap profiling and listen for SIGUSR1 to dump profiles. /// -/// Send `kill -USR1 ` to write a heap profile to `/tmp/moq-relay.heap.`. -pub fn init() { +/// The dump path is controlled by `MALLOC_CONF=prof_prefix:`. +/// Returns `pending` if profiling is not available (i.e. MALLOC_CONF=prof:true was not set). +pub async fn run() -> anyhow::Result<()> { let prof_active = b"prof.active\0"; match unsafe { raw::read::(prof_active) } { @@ -14,28 +15,18 @@ pub fn init() { } Err(err) => { tracing::warn!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable"); - return; + return std::future::pending().await; } } - tokio::spawn(async { - if let Err(err) = signal_handler().await { - tracing::error!(%err, "jemalloc signal handler failed"); - } - }); -} - -async fn signal_handler() -> anyhow::Result<()> { let mut sig = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::user_defined1())?; - let mut seq = 0u64; loop { sig.recv().await; - let path = format!("/tmp/moq-relay.heap.{seq}\0"); - seq += 1; - match unsafe { raw::write(b"prof.dump\0", path.as_ptr()) } { - Ok(()) => tracing::info!(path = &path[..path.len() - 1], "heap profile dumped"), + // Empty path tells jemalloc to use prof_prefix from MALLOC_CONF. + match unsafe { raw::write(b"prof.dump\0", b"\0" as *const u8) } { + Ok(()) => tracing::info!("heap profile dumped"), Err(err) => tracing::error!(%err, "failed to dump heap profile"), } } diff --git a/rs/moq-relay/src/lib.rs b/rs/moq-relay/src/lib.rs index b8fae4285b..ff44df5c9d 100644 --- a/rs/moq-relay/src/lib.rs +++ b/rs/moq-relay/src/lib.rs @@ -12,7 +12,7 @@ mod cluster; mod config; mod connection; #[cfg(feature = "jemalloc")] -pub mod jemalloc_prof; +pub mod jemalloc; mod web; #[cfg(feature = "websocket")] mod websocket; diff --git a/rs/moq-relay/src/main.rs b/rs/moq-relay/src/main.rs index 14d36d0f62..13f7024ff2 100644 --- a/rs/moq-relay/src/main.rs +++ b/rs/moq-relay/src/main.rs @@ -8,9 +8,6 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[tokio::main] async fn main() -> anyhow::Result<()> { - #[cfg(feature = "jemalloc")] - jemalloc_prof::init(); - // TODO: It would be nice to remove this and rely on feature flags only. // However, some dependency is pulling in `ring` and I don't know why, so meh for now. rustls::crypto::aws_lc_rs::default_provider() @@ -55,10 +52,16 @@ async fn main() -> anyhow::Result<()> { // Notify systemd that we're ready after all initialization is complete let _ = sd_notify::notify(&[sd_notify::NotifyState::Ready]); + #[cfg(feature = "jemalloc")] + let jemalloc = jemalloc::run(); + #[cfg(not(feature = "jemalloc"))] + let jemalloc = std::future::pending::>(); + tokio::select! { Err(err) = cluster.clone().run() => return Err(err).context("cluster failed"), Err(err) = web.run() => return Err(err).context("web server failed"), Err(err) = serve(server, cluster, auth) => return Err(err).context("server failed"), + Err(err) = jemalloc => return Err(err).context("jemalloc profiler failed"), else => Ok(()), } } From d4027e01db1405b1188cd4527cd847a0c6fbeb09 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 15:11:54 -0700 Subject: [PATCH 3/6] Update rs/moq-relay/src/jemalloc.rs Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> --- rs/moq-relay/src/jemalloc.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/rs/moq-relay/src/jemalloc.rs b/rs/moq-relay/src/jemalloc.rs index fdc89c81d2..52a3468d10 100644 --- a/rs/moq-relay/src/jemalloc.rs +++ b/rs/moq-relay/src/jemalloc.rs @@ -11,7 +11,10 @@ pub async fn run() -> anyhow::Result<()> { Ok(true) => tracing::info!("jemalloc heap profiling is active"), Ok(false) => { tracing::info!("jemalloc profiling compiled in; activating"); - unsafe { raw::write(prof_active, true) }.ok(); + if let Err(err) = unsafe { raw::write(prof_active, true) } { + tracing::warn!(%err, "failed to activate jemalloc profiling"); + return; + } } Err(err) => { tracing::warn!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable"); From d47d0c0dfef438f296d0f7a2fa572f9794c2d43d Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 15:25:52 -0700 Subject: [PATCH 4/6] Fix sorted dependencies and clean up jemalloc profiling - Sort tikv-jemalloc-ctl and tikv-jemallocator in Cargo.toml (CI fix) - Return error on activation failure instead of pending future - Use debug! instead of warn! when profiling isn't configured - Return Ok(()) when profiling unavailable instead of pending forever - Merge main Co-Authored-By: Claude Opus 4.6 (1M context) --- rs/moq-relay/Cargo.toml | 4 ++-- rs/moq-relay/src/jemalloc.rs | 12 +++++------- 2 files changed, 7 insertions(+), 9 deletions(-) diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index 75a14951bb..d2b69464e2 100644 --- a/rs/moq-relay/Cargo.toml +++ b/rs/moq-relay/Cargo.toml @@ -47,11 +47,11 @@ rustls = { version = "0.23", features = [ serde = { version = "1", features = ["derive"] } serde_with = { version = "3", features = ["json", "base64"] } thiserror = "2" +tikv-jemalloc-ctl = { version = "0.6", optional = true } +tikv-jemallocator = { version = "0.6", features = ["profiling", "unprefixed_malloc_on_supported_platforms"], optional = true } tokio = { workspace = true, features = ["full"] } toml = "0.9" tower-http = { version = "0.6", features = ["cors"] } -tikv-jemallocator = { version = "0.6", features = ["profiling", "unprefixed_malloc_on_supported_platforms"], optional = true } -tikv-jemalloc-ctl = { version = "0.6", optional = true } tracing = "0.1" url = { version = "2", features = ["serde"] } diff --git a/rs/moq-relay/src/jemalloc.rs b/rs/moq-relay/src/jemalloc.rs index 52a3468d10..f792d4a3d8 100644 --- a/rs/moq-relay/src/jemalloc.rs +++ b/rs/moq-relay/src/jemalloc.rs @@ -1,9 +1,10 @@ +use anyhow::Context; use tikv_jemalloc_ctl::raw; /// Activate jemalloc heap profiling and listen for SIGUSR1 to dump profiles. /// /// The dump path is controlled by `MALLOC_CONF=prof_prefix:`. -/// Returns `pending` if profiling is not available (i.e. MALLOC_CONF=prof:true was not set). +/// Returns `Ok(())` if profiling is not available (i.e. MALLOC_CONF=prof:true was not set). pub async fn run() -> anyhow::Result<()> { let prof_active = b"prof.active\0"; @@ -11,14 +12,11 @@ pub async fn run() -> anyhow::Result<()> { Ok(true) => tracing::info!("jemalloc heap profiling is active"), Ok(false) => { tracing::info!("jemalloc profiling compiled in; activating"); - if let Err(err) = unsafe { raw::write(prof_active, true) } { - tracing::warn!(%err, "failed to activate jemalloc profiling"); - return; - } + unsafe { raw::write(prof_active, true) }.context("failed to activate jemalloc profiling")?; } Err(err) => { - tracing::warn!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable"); - return std::future::pending().await; + tracing::debug!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable"); + return Ok(()); } } From 0ee2142968548d2bc24a0038fb41e88bdc155cde Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 15:46:41 -0700 Subject: [PATCH 5/6] Disable Nix fortify hardening for jemalloc builds jemalloc's configure script uses -O0 for test builds, which conflicts with Nix's _FORTIFY_SOURCE hardening (requires optimization). Disable fortify in both the dev shell and the moq-relay package build. Co-Authored-By: Claude Opus 4.6 (1M context) --- flake.nix | 4 ++++ nix/overlay.nix | 3 +++ 2 files changed, 7 insertions(+) diff --git a/flake.nix b/flake.nix index 75aa277af2..72e0852081 100644 --- a/flake.nix +++ b/flake.nix @@ -115,6 +115,10 @@ devShells.default = pkgs.mkShell { packages = rustDeps ++ jsDeps ++ pyDeps ++ cdnDeps; + # jemalloc's configure uses -O0 test builds, which conflict with + # Nix's _FORTIFY_SOURCE hardening (requires -O). + hardeningDisable = [ "fortify" ]; + shellHook = '' export LIBCLANG_PATH="${pkgs.libclang.lib}/lib" ''; diff --git a/nix/overlay.nix b/nix/overlay.nix index abe7d32cc9..aa16a6e2bd 100644 --- a/nix/overlay.nix +++ b/nix/overlay.nix @@ -16,6 +16,9 @@ in # Enable frame pointers for profiling support (negligible overhead on x86_64). # This also ensures the CDN build matches what Cachix caches. RUSTFLAGS = "-C force-frame-pointers=yes"; + # jemalloc's configure uses -O0 test builds, which conflict with + # Nix's _FORTIFY_SOURCE hardening (requires -O). + hardeningDisable = [ "fortify" ]; } ); From 2018ae26089ef64d77588ec5f2ddce2047849d10 Mon Sep 17 00:00:00 2001 From: Luke Curley Date: Thu, 2 Apr 2026 17:15:31 -0700 Subject: [PATCH 6/6] Fix jemalloc profiling build: tikv_jemalloc_ctl::Error doesn't impl std::error::Error Replace anyhow::Context with map_err since tikv_jemalloc_ctl::Error doesn't satisfy the StdError trait bound required by .context(). Also merge latest main to pick up recent changes. Co-Authored-By: Claude Opus 4.6 (1M context) --- rs/moq-relay/src/jemalloc.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rs/moq-relay/src/jemalloc.rs b/rs/moq-relay/src/jemalloc.rs index f792d4a3d8..dbc2b009d0 100644 --- a/rs/moq-relay/src/jemalloc.rs +++ b/rs/moq-relay/src/jemalloc.rs @@ -1,4 +1,3 @@ -use anyhow::Context; use tikv_jemalloc_ctl::raw; /// Activate jemalloc heap profiling and listen for SIGUSR1 to dump profiles. @@ -12,7 +11,8 @@ pub async fn run() -> anyhow::Result<()> { Ok(true) => tracing::info!("jemalloc heap profiling is active"), Ok(false) => { tracing::info!("jemalloc profiling compiled in; activating"); - unsafe { raw::write(prof_active, true) }.context("failed to activate jemalloc profiling")?; + unsafe { raw::write(prof_active, true) } + .map_err(|err| anyhow::anyhow!("failed to activate jemalloc profiling: {err}"))?; } Err(err) => { tracing::debug!(%err, "jemalloc profiling not available — set MALLOC_CONF=prof:true to enable");