diff --git a/Cargo.lock b/Cargo.lock index 1ef854c411..77257a881d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3678,6 +3678,8 @@ dependencies = [ "serde_with", "tempfile", "thiserror 2.0.18", + "tikv-jemalloc-ctl", + "tikv-jemallocator", "tokio", "toml", "tower-http", @@ -6468,6 +6470,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/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/modules/moq-relay.nix b/nix/modules/moq-relay.nix index 2e3b7287d9..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 { @@ -207,6 +213,9 @@ in }; environment = { + # Enable jemalloc heap profiling; dump with `kill -USR1 ` + MALLOC_CONF = "prof:true,prof_active:true,prof_prefix:${cfg.heapDumpPrefix}"; + MOQ_LOG_LEVEL = lib.mkDefault cfg.logLevel; # Server configuration diff --git a/nix/overlay.nix b/nix/overlay.nix index 5bd2ab0161..aa16a6e2bd 100644 --- a/nix/overlay.nix +++ b/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" ]; } ); diff --git a/rs/moq-relay/Cargo.toml b/rs/moq-relay/Cargo.toml index 42f0221327..34f2643c7f 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"] @@ -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"] } diff --git a/rs/moq-relay/src/jemalloc.rs b/rs/moq-relay/src/jemalloc.rs new file mode 100644 index 0000000000..dbc2b009d0 --- /dev/null +++ b/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:`. +/// 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::(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"), + } + } +} diff --git a/rs/moq-relay/src/lib.rs b/rs/moq-relay/src/lib.rs index 455c4fc4cc..ff44df5c9d 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; 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..13f7024ff2 100644 --- a/rs/moq-relay/src/main.rs +++ b/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::>(); + 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(()), } }