Add jemalloc heap profiling to moq-relay

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"
           '';

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 {
@@ -207,6 +213,9 @@ in
       };
 
       environment = {
+        # Enable jemalloc heap profiling; dump with `kill -USR1 <pid>`
+        MALLOC_CONF = "prof:true,prof_active:true,prof_prefix:${cfg.heapDumpPrefix}";
+
         MOQ_LOG_LEVEL = lib.mkDefault cfg.logLevel;
 
         # Server configuration

nix/overlay.nix

@@ -12,10 +12,13 @@ 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";
+      # jemalloc's configure uses -O0 test builds, which conflict with
+      # Nix's _FORTIFY_SOURCE hardening (requires -O).
+      hardeningDisable = [ "fortify" ];
     }
   );
 

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"]
@@ -50,6 +51,8 @@ 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"] }

rs/moq-relay/src/jemalloc.rs

@@ -0,0 +1,34 @@
+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 `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";
+
+	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) }
+				.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");
+			return Ok(());
+		}
+	}
+
+	let mut sig = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::user_defined1())?;
+
+	loop {
+		sig.recv().await;
+
+		// 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"),
+		}
+	}
+}

rs/moq-relay/src/lib.rs

@@ -11,6 +11,8 @@ mod auth;
 mod cluster;
 mod config;
 mod connection;
+#[cfg(feature = "jemalloc")]
+pub mod jemalloc;
 mod web;
 #[cfg(feature = "websocket")]
 mod websocket;

rs/moq-relay/src/main.rs

@@ -2,6 +2,10 @@ 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<()> {
 	// TODO: It would be nice to remove this and rely on feature flags only.
@@ -48,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(()),
 	}
 }