diff --git a/Cargo.lock b/Cargo.lock index 3c5c2e985..e78c6ea07 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -444,6 +444,7 @@ dependencies = [ "rusqlite", "serde", "serde_json", + "tikv-jemallocator", "tokio", "tokio-postgres", "tracing", @@ -492,7 +493,7 @@ dependencies = [ "postgres-types", "proc-macro2", "prometheus-http-query", - "prost", + "prost 0.12.6", "quote", "rand 0.8.5", "rand_chacha 0.3.1", @@ -535,7 +536,7 @@ dependencies = [ "arroyo-types", "base64 0.21.7", "dlopen2", - "prost", + "prost 0.12.6", "serde_json", "tokio", "toml", @@ -578,7 +579,7 @@ dependencies = [ "object_store", "once_cell", "parquet", - "prost", + "prost 0.12.6", "rand 0.8.5", "rdkafka", "rdkafka-sys", @@ -636,7 +637,7 @@ dependencies = [ "prettyplease 0.2.20", "proc-macro2", "prometheus", - "prost", + "prost 0.12.6", "quote", "rand 0.8.5", "refinery", @@ -675,7 +676,7 @@ dependencies = [ "hex", "petgraph", "proc-macro2", - "prost", + "prost 0.12.6", "quote", "rand 0.8.5", "regex", @@ -715,7 +716,7 @@ dependencies = [ "petgraph", "prettyplease 0.2.20", "proc-macro2", - "prost", + "prost 0.12.6", "quote", "regex", "rstest", @@ -748,7 +749,7 @@ dependencies = [ "bincode", "chrono", "memchr", - "prost", + "prost 0.12.6", "schemars", "serde", "serde_json", @@ -780,7 +781,7 @@ dependencies = [ "lazy_static", "local-ip-address", "prometheus", - "prost", + "prost 0.12.6", "rand 0.8.5", "tokio", "tokio-stream", @@ -829,7 +830,7 @@ dependencies = [ "datafusion", "dlopen2", "futures", - "prost", + "prost 0.12.6", "rand 0.8.5", "serde", "serde_json", @@ -860,7 +861,7 @@ dependencies = [ "k8s-openapi", "log", "nanoid", - "prost", + "prost 0.12.6", "rand 0.8.5", "regex", "reqwest 0.11.27", @@ -888,6 +889,7 @@ dependencies = [ "dirs", "futures", "hyper 0.14.30", + "jemalloc_pprof", "lazy_static", "once_cell", "prometheus", @@ -973,7 +975,7 @@ dependencies = [ "once_cell", "parquet", "prometheus", - "prost", + "prost 0.12.6", "rand 0.8.5", "tokio", "tonic", @@ -1116,7 +1118,7 @@ dependencies = [ "ordered-float 3.9.2", "petgraph", "prometheus", - "prost", + "prost 0.12.6", "rand 0.8.5", "regex", "regress 0.6.0", @@ -3297,7 +3299,7 @@ dependencies = [ "datafusion-expr", "datafusion-proto-common", "object_store", - "prost", + "prost 0.12.6", ] [[package]] @@ -3309,7 +3311,7 @@ dependencies = [ "chrono", "datafusion-common", "object_store", - "prost", + "prost 0.12.6", ] [[package]] @@ -5116,6 +5118,15 @@ version = "1.70.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + [[package]] name = "itertools" version = "0.11.0" @@ -5149,6 +5160,23 @@ version = "1.0.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b" +[[package]] +name = "jemalloc_pprof" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96368c0fc161a0a1a20b3952b6fd31ee342fffc87ed9e48ac1ed49fb25686655" +dependencies = [ + "anyhow", + "libc", + "mappings", + "once_cell", + "pprof_util", + "tempfile", + "tikv-jemalloc-ctl", + "tokio", + "tracing", +] + [[package]] name = "jobserver" version = "0.1.32" @@ -5613,6 +5641,19 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" +[[package]] +name = "mappings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fa2605f461115ef6336342b12f0d8cabdfd7b258fed86f5f98c725535843601" +dependencies = [ + "anyhow", + "libc", + "once_cell", + "pprof_util", + "tracing", +] + [[package]] name = "matchers" version = "0.1.0" @@ -6630,6 +6671,19 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" +[[package]] +name = "pprof_util" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c620a1858d6ebf10d7c60256629078b2d106968d0e6ff63b850d9ecd84008fbe" +dependencies = [ + "anyhow", + "flate2", + "num", + "paste", + "prost 0.11.9", +] + [[package]] name = "ppv-lite86" version = "0.2.17" @@ -6837,6 +6891,16 @@ dependencies = [ "url", ] +[[package]] +name = "prost" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b82eaa1d779e9a4bc1c3217db8ffbeabaae1dca241bf70183242128d48681cd" +dependencies = [ + "bytes", + "prost-derive 0.11.9", +] + [[package]] name = "prost" version = "0.12.6" @@ -6844,7 +6908,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "deb1435c188b76130da55f17a466d252ff7b1418b2ad3e037d127b94e3411f29" dependencies = [ "bytes", - "prost-derive", + "prost-derive 0.12.6", ] [[package]] @@ -6861,13 +6925,26 @@ dependencies = [ "once_cell", "petgraph", "prettyplease 0.2.20", - "prost", + "prost 0.12.6", "prost-types", "regex", "syn 2.0.72", "tempfile", ] +[[package]] +name = "prost-derive" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5d2d8d10f3c6ded6da8b05b5fb3b8a5082514344d56c9f871412d29b4e075b4" +dependencies = [ + "anyhow", + "itertools 0.10.5", + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "prost-derive" version = "0.12.6" @@ -6887,7 +6964,7 @@ version = "0.12.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9091c90b0a32608e984ff2fa4091273cbdd755d54935c51d520887f4a1dbd5b0" dependencies = [ - "prost", + "prost 0.12.6", ] [[package]] @@ -8626,6 +8703,37 @@ dependencies = [ "ordered-float 2.10.1", ] +[[package]] +name = "tikv-jemalloc-ctl" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "619bfed27d807b54f7f776b9430d4f8060e66ee138a28632ca898584d462c31c" +dependencies = [ + "libc", + "paste", + "tikv-jemalloc-sys", +] + +[[package]] +name = "tikv-jemalloc-sys" +version = "0.5.4+5.3.0-patched" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9402443cb8fd499b6f327e40565234ff34dbda27460c5b47db0db77443dd85d1" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "tikv-jemallocator" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "965fe0c26be5c56c94e38ba547249074803efd52adfb66de62107d95aab3eaca" +dependencies = [ + "libc", + "tikv-jemalloc-sys", +] + [[package]] name = "time" version = "0.3.36" @@ -8896,7 +9004,7 @@ dependencies = [ "hyper-timeout 0.4.1", "percent-encoding", "pin-project", - "prost", + "prost 0.12.6", "tokio", "tokio-stream", "tower", @@ -8925,7 +9033,7 @@ version = "0.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "548c227bd5c0fae5925812c4ec6c66ffcfced23ea370cb823f4d18f0fc1cb6a7" dependencies = [ - "prost", + "prost 0.12.6", "prost-types", "tokio", "tokio-stream", diff --git a/crates/arroyo-server-common/Cargo.toml b/crates/arroyo-server-common/Cargo.toml index 6215a519e..49ce83f14 100644 --- a/crates/arroyo-server-common/Cargo.toml +++ b/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" diff --git a/crates/arroyo-server-common/src/lib.rs b/crates/arroyo-server-common/src/lib.rs index 72899cc1e..b6bef152e 100644 --- a/crates/arroyo-server-common/src/lib.rs +++ b/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>) -> String { .unwrap() } + +pub async fn handle_get_heap() -> Result { + 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); diff --git a/crates/arroyo/Cargo.toml b/crates/arroyo/Cargo.toml index a81976d92..749fdd783 100644 --- a/crates/arroyo/Cargo.toml +++ b/crates/arroyo/Cargo.toml @@ -3,6 +3,9 @@ name = "arroyo" version = "0.12.0-dev" edition = "2021" +[features] +profiling = ["tikv-jemallocator/profiling"] + [dependencies] arroyo-types = { path ="../arroyo-types" } arroyo-connectors = { path ="../arroyo-connectors" } @@ -36,3 +39,5 @@ rand = "0.8.5" reqwest = "0.11" clio = { version = "0.3.5", features = ["clap", "clap-parse"] } async-trait = "0.1.80" + +tikv-jemallocator = { version = "0.5.4", features = ["unprefixed_malloc_on_supported_platforms"]} \ No newline at end of file diff --git a/crates/arroyo/src/main.rs b/crates/arroyo/src/main.rs index ae112b198..3620cb0b0 100644 --- a/crates/arroyo/src/main.rs +++ b/crates/arroyo/src/main.rs @@ -23,6 +23,16 @@ use tokio_postgres::{Client, Connection, NoTls}; use tracing::{debug, error, info}; use uuid::Uuid; +#[cfg(not(target_env = "msvc"))] +#[global_allocator] +static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; + +#[cfg(feature = "profiling")] +#[allow(non_upper_case_globals)] +#[export_name = "malloc_conf"] +pub static malloc_conf: &[u8] = b"prof:true,prof_active:true,lg_prof_sample:19\0"; + + #[derive(Parser)] #[command(version, about)] struct Cli {