ROX-30302: implement heap profiler

Cargo.toml

@@ -25,6 +25,8 @@ pprof = { version = "0.15.0", features = ["prost-codec"] }
 prometheus-client = { version = "0.24.0", default-features = false }
 prost = "0.13.5"
 prost-types = "0.13.5"
+tikv-jemalloc-ctl = { version = "0.6.0", features = ["profiling"] }
+tikv-jemallocator = { version = "0.6.0", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] }
 serde = { version = "1.0.219", features = ["derive"] }
 serde_json = "1.0.142"
 tokio = { version = "1.40.0", default-features = false, features = [

fact/Cargo.toml

@@ -22,6 +22,9 @@ pprof = { workspace = true }
 prometheus-client = { workspace = true }
 prost = { workspace = true }
 prost-types = { workspace = true }
+tempfile = { workspace = true, optional = true}
+tikv-jemalloc-ctl = { workspace = true, optional = true }
+tikv-jemallocator = { workspace = true, optional = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 uuid = { workspace = true }
@@ -42,3 +45,4 @@ path = "src/main.rs"
 
 [features]
 bpf-test = []
+jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl", "tempfile"]

fact/src/endpoints.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "jemalloc")]
+use std::ffi::CString;
 use std::{future::Future, net::SocketAddr, pin::Pin};
 
 use http_body_util::{BodyExt, Full};
@@ -9,6 +11,10 @@ use hyper::{
 };
 use hyper_util::rt::TokioIo;
 use log::{info, warn};
+#[cfg(feature = "jemalloc")]
+use tempfile::NamedTempFile;
+#[cfg(feature = "jemalloc")]
+use tikv_jemalloc_ctl::raw as mallctl;
 use tokio::{net::TcpListener, sync::watch, task::JoinHandle};
 
 use crate::metrics::exporter::Exporter;
@@ -148,7 +154,7 @@ impl Server {
                 Ok(_) => Server::response_with_content_type(
                     StatusCode::OK,
                     "text/plain",
-                    "CPU profiler starter",
+                    "CPU profiler started",
                 ),
                 Err(e) => Server::response(
                     StatusCode::INTERNAL_SERVER_ERROR,
@@ -190,6 +196,79 @@ impl Server {
             ),
         }
     }
+
+    #[cfg(feature = "jemalloc")]
+    async fn handle_heap_profiler(&self, body: Incoming) -> ServerResponse {
+        let body = match body.collect().await {
+            Ok(b) => b.to_bytes(),
+            Err(e) => {
+                return Server::response(
+                    StatusCode::BAD_REQUEST,
+                    format!("Failed to read request body: {e}"),
+                )
+            }
+        };
+
+        let state = if body == "on" {
+            true
+        } else if body == "off" {
+            false
+        } else {
+            return Server::response(
+                StatusCode::BAD_REQUEST,
+                format!("Invalid request body: {body:?}"),
+            );
+        };
+
+        let res = unsafe { mallctl::update(b"prof.active\0", true) };
+
+        match res {
+            Ok(_) => Server::response_with_content_type(
+                StatusCode::OK,
+                "text/plain",
+                format!(
+                    "Heap profiler {}",
+                    if state { "started" } else { "stopped" }
+                ),
+            ),
+            Err(e) => Server::response(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!(
+                    "Failed to {} heap profiler: {e}",
+                    if state { "start" } else { "stop" }
+                ),
+            ),
+        }
+    }
+
+    #[cfg(feature = "jemalloc")]
+    async fn handle_heap_report(&self) -> ServerResponse {
+        let f = match NamedTempFile::new() {
+            Ok(f) => f,
+            Err(e) => {
+                return Server::response(
+                    StatusCode::INTERNAL_SERVER_ERROR,
+                    format!("Failed to create dump file: {e}"),
+                );
+            }
+        };
+        let path = CString::new(f.path().as_os_str().as_encoded_bytes()).unwrap();
+
+        if let Err(e) = unsafe { mallctl::write(b"prof.dump\0", path.as_ptr()) } {
+            return Server::response(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Failed to dump heap profile: {e}"),
+            );
+        }
+
+        match std::fs::read_to_string(f.path()) {
+            Ok(profile) => Server::response(StatusCode::OK, profile),
+            Err(e) => Server::response(
+                StatusCode::INTERNAL_SERVER_ERROR,
+                format!("Failed to read heap profile dump: {e}"),
+            ),
+        }
+    }
 }
 
 impl Service<Request<Incoming>> for Server {
@@ -205,6 +284,15 @@ impl Service<Request<Incoming>> for Server {
                 (&Method::GET, "/health_check") => s.handle_health_check(),
                 (&Method::POST, "/profile/cpu") => s.handle_cpu_profiler(req.into_body()).await,
                 (&Method::GET, "/profile/cpu") => s.handle_cpu_report().await,
+                #[cfg(feature = "jemalloc")]
+                (&Method::POST, "/profile/heap") => s.handle_heap_profiler(req.into_body()).await,
+                #[cfg(feature = "jemalloc")]
+                (&Method::GET, "/profile/heap") => s.handle_heap_report().await,
+                #[cfg(not(feature = "jemalloc"))]
+                (_, "/profile/heap") => Server::response(
+                    StatusCode::SERVICE_UNAVAILABLE,
+                    "Heap profiler not supported",
+                ),
                 (&Method::GET, "/profile") => s.handle_profiler_status().await,
                 _ => Server::response(StatusCode::NOT_FOUND, ""),
             }

fact/src/main.rs

@@ -1,5 +1,13 @@
 use fact::config::FactConfig;
 
+#[cfg(feature = "jemalloc")]
+#[global_allocator]
+static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
+#[allow(non_upper_case_globals)]
+#[export_name = "malloc_conf"]
+pub static malloc_conf: &[u8] = b"prof:true,prof_active:false\0";
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     fact::init_log()?;