Add pprof heap endpoints to node api

Cargo.toml

@@ -169,8 +169,7 @@ tempfile = "3.6.0"
 test-log = { version = "0.2.11", default-features = false, features = [
     "trace",
 ] }
-# tikv-jemallocator has not yet been released with musl target support, so we pin a main commit
-tikv-jemallocator = { git = "https://github.com/restatedev/jemallocator", rev = "7c32f6e3d6ad5e4e492cc08d6bdb8307acf9afa0", default-features = false }
+tikv-jemallocator = "0.6"
 thiserror = "1.0"
 tokio = { version = "1.40.0", default-features = false, features = [
     "rt-multi-thread",

crates/node/Cargo.toml

@@ -72,6 +72,9 @@ tower = { workspace = true }
 tower-http = { workspace = true, features = ["trace"] }
 tracing = { workspace = true }
 
+[target.'cfg(not(target_env = "msvc"))'.dependencies]
+jemalloc_pprof = "0.6.0"
+
 [dev-dependencies]
 restate-test-util = { workspace = true }
 

crates/node/src/network_server/mod.rs

@@ -11,6 +11,7 @@
 mod handler;
 mod metrics;
 mod multiplex;
+mod pprof;
 mod prometheus_helpers;
 mod service;
 mod state;

crates/node/src/network_server/pprof.rs

@@ -0,0 +1,82 @@
+#[cfg(target_os = "linux")]
+mod pprof {
+    use http::StatusCode;
+
+    pub async fn heap() -> Result<impl axum::response::IntoResponse, (StatusCode, String)> {
+        match jemalloc_pprof::PROF_CTL.as_ref() {
+            Some(prof_ctl) => {
+                let mut prof_ctl = prof_ctl.lock().await;
+                if prof_ctl.activated() {
+                    let pprof = prof_ctl
+                        .dump_pprof()
+                        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
+                    Ok(pprof)
+                } else {
+                    Err((
+                        axum::http::StatusCode::PRECONDITION_FAILED,
+                        "Heap profiling not activated: first call /debug/pprof/heap/activate\n"
+                            .into(),
+                    ))
+                }
+            }
+            None => Err((
+                axum::http::StatusCode::PRECONDITION_FAILED,
+                "Heap profiling not enabled: run with MALLOC_CONF=\"prof:true\"\n".into(),
+            )),
+        }
+    }
+
+    pub async fn activate_heap() -> Result<(), (StatusCode, String)> {
+        match jemalloc_pprof::PROF_CTL.as_ref() {
+            Some(prof_ctl) => {
+                let mut prof_ctl = prof_ctl.lock().await;
+                prof_ctl
+                    .activate()
+                    .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
+                Ok(())
+            }
+            None => Err((
+                axum::http::StatusCode::PRECONDITION_FAILED,
+                "Heap profiling not enabled: run with MALLOC_CONF=\"prof:true\"\n".into(),
+            )),
+        }
+    }
+
+    pub async fn deactivate_heap() -> Result<(), (StatusCode, String)> {
+        match jemalloc_pprof::PROF_CTL.as_ref() {
+            Some(prof_ctl) => {
+                let mut prof_ctl = prof_ctl.lock().await;
+                prof_ctl
+                    .deactivate()
+                    .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
+                Ok(())
+            }
+            None => Err((
+                axum::http::StatusCode::PRECONDITION_FAILED,
+                "Heap profiling not enabled: run with MALLOC_CONF=\"prof:true\"\n".into(),
+            )),
+        }
+    }
+}
+
+#[cfg(not(target_os = "linux"))]
+mod pprof {
+    use http::StatusCode;
+
+    pub async fn heap() -> (StatusCode, String) {
+        (
+            axum::http::StatusCode::PRECONDITION_FAILED,
+            "Heap profiling is only available on Linux\n".into(),
+        )
+    }
+
+    pub async fn activate_heap() -> (StatusCode, String) {
+        heap().await
+    }
+
+    pub async fn deactivate_heap() -> (StatusCode, String) {
+        heap().await
+    }
+}
+
+pub use pprof::*;

crates/node/src/network_server/service.rs

@@ -28,6 +28,7 @@ use restate_metadata_store::MetadataStoreClient;
 use restate_storage_query_datafusion::context::QueryContext;
 use restate_types::config::CommonOptions;
 
+use super::pprof;
 use crate::network_server::handler;
 use crate::network_server::handler::cluster_ctrl::ClusterCtrlSvcHandler;
 use crate::network_server::handler::node::NodeSvcHandler;
@@ -40,7 +41,6 @@ pub struct NetworkServer {
     worker_deps: Option<WorkerDependencies>,
     admin_deps: Option<AdminDependencies>,
 }
-
 impl NetworkServer {
     pub fn new(
         connection_manager: ConnectionManager,
@@ -74,6 +74,9 @@ impl NetworkServer {
         // -- HTTP service (for prometheus et al.)
         let router = axum::Router::new()
             .route("/metrics", get(handler::render_metrics))
+            .route("/debug/pprof/heap", get(pprof::heap))
+            .route("/debug/pprof/heap/activate", get(pprof::activate_heap))
+            .route("/debug/pprof/heap/deactivate", get(pprof::deactivate_heap))
             .with_state(shared_state)
             .layer(TraceLayer::new_for_http().make_span_with(span_factory.clone()))
             .fallback(handler_404);

server/src/main.rs

@@ -43,6 +43,12 @@ use tikv_jemallocator::Jemalloc;
 #[global_allocator]
 static GLOBAL: Jemalloc = Jemalloc;
 
+// On linux, run jemalloc with profiling enabled, but inactive (so there is no performance impact)
+// If needed, profiling can be activated with :5122/debug/pprof/heap/activate, or by overriding this value with $MALLOC_CONF
+#[cfg(target_os = "linux")]
+#[export_name = "malloc_conf"]
+pub static MALLOC_CONF: &[u8] = b"prof:true,prof_active:false,lg_prof_sample:19\0";
+
 #[derive(Debug, clap::Parser)]
 #[command(author, version, about)]
 struct RestateArguments {