experiment: expose ability to capture jemalloc heap profiles in pprof format from admin API

.cargo/config.toml

@@ -2,4 +2,4 @@
 rustflags = ["--cfg", "tokio_unstable"]
 
 [env]
-JEMALLOC_SYS_WITH_MALLOC_CONF = "abort_conf:true,max_background_threads:1,narenas:1,tcache:false,thp:never,oversize_threshold:32768,dirty_decay_ms:1000,muzzy_decay_ms:0"
+JEMALLOC_SYS_WITH_MALLOC_CONF = "abort_conf:true,prof:true,prof_active:true,lg_prof_sample:19,max_background_threads:1,narenas:1,tcache:false,thp:never,oversize_threshold:32768,dirty_decay_ms:1000,muzzy_decay_ms:0"

Cargo.toml

@@ -66,6 +66,7 @@ tracing = { version = "0.1", default-features = false, features = ["std"] }
 ahash = { version = "0.8", default-features = false, features = ["std", "runtime-rng"] }
 async-compression = { version = "0.4.13", default-features = false }
 bitmask-enum = { version = "2.2", default-features = false }
+jemalloc_pprof = { version = "0.6", default-features = false }
 figment = { version = "0.10", default-features = false }
 hostname = { version = "0.4", default-features = false }
 http = { version = "1", default-features = false }

LICENSE-3rdparty.csv

@@ -140,6 +140,7 @@ is-terminal,https://github.com/sunfishcode/is-terminal,MIT,"softprops <d.tangren
 is_terminal_polyfill,https://github.com/polyfill-rs/is_terminal_polyfill,MIT OR Apache-2.0,The is_terminal_polyfill Authors
 itertools,https://github.com/rust-itertools/itertools,MIT OR Apache-2.0,bluss
 itoa,https://github.com/dtolnay/itoa,MIT OR Apache-2.0,David Tolnay <[email protected]>
+jemalloc_pprof,https://github.com/polarsignals/rust-jemalloc-pprof,Apache-2.0,"Frederic Branczyk <[email protected]>, Brennan Vincent <[email protected]>"
 jobserver,https://github.com/rust-lang/jobserver-rs,MIT OR Apache-2.0,Alex Crichton <[email protected]>
 js-sys,https://github.com/rustwasm/wasm-bindgen/tree/master/crates/js-sys,MIT OR Apache-2.0,The wasm-bindgen Developers
 jsonpath-rust,https://github.com/besok/jsonpath-rust,MIT,BorisZhguchev <[email protected]>
@@ -154,6 +155,7 @@ libm,https://github.com/rust-lang/libm,MIT AND (MIT OR Apache-2.0),Jorge Aparici
 linux-raw-sys,https://github.com/sunfishcode/linux-raw-sys,Apache-2.0 WITH LLVM-exception OR Apache-2.0 OR MIT,Dan Gohman <[email protected]>
 litemap,https://github.com/unicode-org/icu4x,Unicode-3.0,The ICU4X Project Developers
 log,https://github.com/rust-lang/log,MIT OR Apache-2.0,The Rust Project Developers
+mappings,https://github.com/polarsignals/rust-jemalloc-pprof,Apache-2.0,The mappings Authors
 matchers,https://github.com/hawkw/matchers,MIT,Eliza Weisman <[email protected]>
 matchit,https://github.com/ibraheemdev/matchit,MIT AND BSD-3-Clause,Ibraheem Ahmed <[email protected]>
 matrixmultiply,https://github.com/bluss/matrixmultiply,MIT OR Apache-2.0,"bluss, R. Janis Goldschmidt"
@@ -169,9 +171,13 @@ ndarray,https://github.com/rust-ndarray/ndarray,MIT OR Apache-2.0,"Ulrik Sverdru
 noisy_float,https://github.com/SergiusIW/noisy_float-rs,Apache-2.0,Matthew Michelotti <[email protected]>
 nom,https://github.com/Geal/nom,MIT,[email protected]
 nu-ansi-term,https://github.com/nushell/nu-ansi-term,MIT,"[email protected], Ryan Scheel (Havvy) <[email protected]>, Josh Triplett <[email protected]>, The Nushell Project Developers"
+num,https://github.com/rust-num/num,MIT OR Apache-2.0,The Rust Project Developers
+num-bigint,https://github.com/rust-num/num-bigint,MIT OR Apache-2.0,The Rust Project Developers
 num-complex,https://github.com/rust-num/num-complex,MIT OR Apache-2.0,The Rust Project Developers
 num-conv,https://github.com/jhpratt/num-conv,MIT OR Apache-2.0,Jacob Pratt <[email protected]>
 num-integer,https://github.com/rust-num/num-integer,MIT OR Apache-2.0,The Rust Project Developers
+num-iter,https://github.com/rust-num/num-iter,MIT OR Apache-2.0,The Rust Project Developers
+num-rational,https://github.com/rust-num/num-rational,MIT OR Apache-2.0,The Rust Project Developers
 num-traits,https://github.com/rust-num/num-traits,MIT OR Apache-2.0,The Rust Project Developers
 num_threads,https://github.com/jhpratt/num_threads,MIT OR Apache-2.0,Jacob Pratt <[email protected]>
 object,https://github.com/gimli-rs/object,Apache-2.0 OR MIT,The object Authors
@@ -202,6 +208,7 @@ plotters,https://github.com/plotters-rs/plotters,MIT,Hao Hou <[email protected]
 portable-atomic,https://github.com/taiki-e/portable-atomic,Apache-2.0 OR MIT,The portable-atomic Authors
 portable-atomic-util,https://github.com/taiki-e/portable-atomic,Apache-2.0 OR MIT,The portable-atomic-util Authors
 powerfmt,https://github.com/jhpratt/powerfmt,MIT OR Apache-2.0,Jacob Pratt <[email protected]>
+pprof_util,https://github.com/polarsignals/rust-jemalloc-pprof,Apache-2.0,The pprof_util Authors
 ppv-lite86,https://github.com/cryptocorrosion/cryptocorrosion,MIT OR Apache-2.0,The CryptoCorrosion Contributors
 prettyplease,https://github.com/dtolnay/prettyplease,MIT OR Apache-2.0,David Tolnay <[email protected]>
 proc-macro2,https://github.com/dtolnay/proc-macro2,MIT OR Apache-2.0,"David Tolnay <[email protected]>, Alex Crichton <[email protected]>"

bin/agent-data-plane/Cargo.toml

@@ -12,6 +12,7 @@ fips = ["saluki-app/tls-fips"]
 [dependencies]
 async-trait = { workspace = true }
 bytesize = { workspace = true }
+jemalloc_pprof = { workspace = true }
 memory-accounting = { workspace = true }
 saluki-app = { workspace = true, features = ["full"] }
 saluki-components = { workspace = true }
@@ -27,19 +28,14 @@ saluki-metadata = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 stringtheory = { workspace = true }
-tokio = { workspace = true, features = [
-  "macros",
-  "rt",
-  "rt-multi-thread",
-  "signal",
-] }
-tokio-rustls = { workspace = true }
+tokio = { workspace = true, features = ["macros", "rt", "rt-multi-thread", "signal"] }
 tracing = { workspace = true }
 
 [target.'cfg(target_os = "linux")'.dependencies]
 tikv-jemalloc-ctl = { workspace = true, features = ["use_std"] }
 tikv-jemallocator = { workspace = true, features = [
   "background_threads",
+  "profiling",
   "unprefixed_malloc_on_supported_platforms",
   "stats",
 ] }

bin/agent-data-plane/src/main.rs

@@ -11,7 +11,7 @@ use std::{
 };
 
 use memory_accounting::{ComponentBounds, ComponentRegistry};
-use saluki_app::{api::APIBuilder, logging::LoggingAPIHandler, prelude::*};
+use saluki_app::{api::APIBuilder, logging::LoggingAPIHandler, memory::MemoryProfilingAPIHandler, prelude::*};
 use saluki_components::{
     destinations::{
         new_remote_agent_service, DatadogEventsServiceChecksConfiguration, DatadogMetricsConfiguration,
@@ -123,6 +123,7 @@ async fn run(started: Instant, logging_api_handler: LoggingAPIHandler) -> Result
         .with_self_signed_tls()
         .with_grpc_service(new_remote_agent_service())
         .with_handler(logging_api_handler)
+        .with_handler(MemoryProfilingAPIHandler)
         .with_optional_handler(env_provider.workload_api_handler());
 
     // Run memory bounds validation to ensure that we can launch the topology with our configured memory limit, if any.

lib/saluki-app/Cargo.toml

@@ -10,7 +10,7 @@ default = []
 full = ["api", "logging", "memory", "metrics", "tls"]
 api = ["dep:axum", "dep:saluki-api", "dep:saluki-error", "dep:saluki-io", "dep:tokio", "dep:tower", "dep:tracing"]
 logging = ["api", "dep:chrono", "dep:chrono-tz", "dep:iana-time-zone", "dep:serde", "dep:tracing", "dep:tracing-subscriber"]
-memory = ["metrics", "dep:bytesize", "dep:memory-accounting", "dep:saluki-config", "dep:saluki-error", "dep:serde", "dep:tokio", "dep:tracing"]
+memory = ["api", "metrics", "dep:bytesize", "dep:jemalloc_pprof", "dep:memory-accounting", "dep:saluki-config", "dep:saluki-error", "dep:serde", "dep:tokio", "dep:tracing"]
 metrics = ["dep:saluki-core", "dep:metrics", "dep:tokio"]
 tls = ["dep:saluki-error", "dep:saluki-tls"]
 tls-fips = ["saluki-tls?/fips"]
@@ -22,6 +22,7 @@ chrono = { workspace = true, optional = true }
 chrono-tz = { workspace = true, optional = true }
 http = { workspace = true }
 iana-time-zone = { workspace = true, optional = true }
+jemalloc_pprof = { workspace = true, default-features = false, optional = true }
 memory-accounting = { workspace = true, optional = true }
 metrics = { workspace = true, optional = true }
 metrics-util = { workspace = true, features = ["handles", "recency", "registry"], optional = true }

lib/saluki-app/src/memory.rs

@@ -13,6 +13,11 @@ use memory_accounting::{
     ComponentRegistry, MemoryGrant, MemoryLimiter, VerifiedBounds,
 };
 use metrics::{counter, gauge, Counter, Gauge, Level};
+use saluki_api::{
+    response::IntoResponse,
+    routing::{get, Router},
+    APIHandler, StatusCode,
+};
 use saluki_config::GenericConfiguration;
 use saluki_error::{generic_error, ErrorContext as _, GenericError};
 use serde::Deserialize;
@@ -316,3 +321,39 @@ impl CgroupMemoryParser {
         memory.parse::<ByteSize>().ok()
     }
 }
+
+/// An API handler for memory profiling.
+///
+/// This handler exposes a single route -- `/debug/pprof/heap` -- which returns a jemalloc heap profile in pprof
+/// format. As one might expect, this handler should only be used when jemalloc is set as the global allocator for the
+/// application.
+pub struct MemoryProfilingAPIHandler;
+
+impl MemoryProfilingAPIHandler {
+    async fn pprof_handler() -> 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)
+    }
+}
+
+impl APIHandler for MemoryProfilingAPIHandler {
+    type State = ();
+
+    fn generate_initial_state(&self) -> Self::State {}
+
+    fn generate_routes(&self) -> Router<Self::State> {
+        Router::new().route("/debug/pprof/heap", get(Self::pprof_handler))
+    }
+}
+
+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()))
+    }
+}