From d82cd273a630559de98f9adfa7adb7425e4e5356 Mon Sep 17 00:00:00 2001
From: Luke Curley <kixelated@gmail.com>
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.<n> 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) <noreply@anthropic.com>
---
 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 <pid>`
+        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 <pid>` to write a heap profile to `/tmp/moq-relay.heap.<seq>`.
+pub fn init() {
+	let prof_active = b"prof.active\0";
+
+	match unsafe { raw::read::<bool>(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 <kixelated@gmail.com>
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) <noreply@anthropic.com>
---
 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 <pid>`
-        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 <pid>` to write a heap profile to `/tmp/moq-relay.heap.<seq>`.
-pub fn init() {
+/// The dump path is controlled by `MALLOC_CONF=prof_prefix:<path>`.
+/// 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::<bool>(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::<anyhow::Result<()>>();
+
 	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 <kixelated@gmail.com>
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 <kixelated@gmail.com>
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) <noreply@anthropic.com>
---
 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:<path>`.
-/// 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 <kixelated@gmail.com>
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) <noreply@anthropic.com>
---
 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 <kixelated@gmail.com>
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) <noreply@anthropic.com>
---
 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");
