Switch to jemalloc

crates/arroyo-server-common/Cargo.toml

@@ -7,6 +7,9 @@ edition = "2021"
 arroyo-types = { path = "../arroyo-types" }
 arroyo-rpc = { path = "../arroyo-rpc" }
 
+# tikv-jemallocator = { version = "0.6.0", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }
+jemalloc_pprof = "0.4.2"
+
 # logging
 tracing = "0.1"
 tracing-logfmt = "0.2.0"

crates/arroyo-server-common/src/lib.rs

@@ -22,6 +22,7 @@ use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::task::{Context, Poll};
+use axum::response::IntoResponse;
 use tonic::body::BoxBody;
 use tonic::transport::Server;
 use tower::layer::util::Stack;
@@ -232,6 +233,25 @@ async fn details<'a>(State(state): State<Arc<AdminState>>) -> String {
     .unwrap()
 }
 
+
+pub async fn handle_get_heap() -> Result<impl IntoResponse, (StatusCode, String)> {
+    let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await;
+    require_profiling_activated(&prof_ctl)?;
+    let pprof = prof_ctl
+        .dump_pprof()
+        .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?;
+    Ok(pprof)
+}
+
+/// Checks whether jemalloc profiling is activated an returns an error response if not.
+fn require_profiling_activated(prof_ctl: &jemalloc_pprof::JemallocProfCtl) -> Result<(), (StatusCode, String)> {
+    if prof_ctl.activated() {
+        Ok(())
+    } else {
+        Err((axum::http::StatusCode::FORBIDDEN, "heap profiling not activated".into()))
+    }
+}
+
 pub async fn start_admin_server(service: &str) -> anyhow::Result<()> {
     let addr = config().admin.bind_address;
     let port = config().admin.http_port;
@@ -248,6 +268,7 @@ pub async fn start_admin_server(service: &str) -> anyhow::Result<()> {
         .route("/metrics.pb", get(metrics_proto))
         .route("/details", get(details))
         .route("/config", get(config_route))
+        .route("/debug/pprof/heap", get(handle_get_heap))
         .with_state(state);
 
     let addr = SocketAddr::new(addr, port);