diff --git a/Cargo.lock b/Cargo.lock index c8727cd4..7341d079 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -958,6 +958,7 @@ dependencies = [ "hyper-rustls 0.26.0", "hyper-util", "indicatif", + "jemalloc_pprof", "lazy_static", "leaky-bucket", "libc", @@ -975,6 +976,7 @@ dependencies = [ "serde_json", "sysinfo", "termion", + "tikv-jemallocator", "tokio", "tokio-rustls 0.25.0", "tokio-stream", @@ -2029,6 +2031,25 @@ version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c" +[[package]] +name = "jemalloc_pprof" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45b38a2cc3eb7b0e332c6368a6fd6a1a603a5be9526f0810f8e0682513538541" +dependencies = [ + "anyhow", + "flate2", + "libc", + "num", + "once_cell", + "paste", + "prost 0.11.9", + "tempfile", + "tikv-jemalloc-ctl", + "tokio", + "tracing", +] + [[package]] name = "jobserver" version = "0.1.27" @@ -2418,6 +2439,20 @@ dependencies = [ "winapi", ] +[[package]] +name = "num" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3135b08af27d103b0a51f2ae0f8632117b7b185ccf931445affa8df530576a41" +dependencies = [ + "num-bigint", + "num-complex", + "num-integer", + "num-iter", + "num-rational", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.4.4" @@ -2429,6 +2464,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23c6602fda94a57c990fe0df199a035d83576b496aa29f4e634a8ac6004e68a6" +dependencies = [ + "num-traits", +] + [[package]] name = "num-format" version = "0.4.4" @@ -2441,19 +2485,41 @@ dependencies = [ [[package]] name = "num-integer" -version = "0.1.45" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.44" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +checksum = "d869c01cc0c455284163fd0092f1f93835385ccab5a98a0dcc497b2f8bf055a9" dependencies = [ "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-rational" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0" +dependencies = [ + "autocfg", + "num-bigint", + "num-integer", "num-traits", ] [[package]] name = "num-traits" -version = "0.2.17" +version = "0.2.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "39e3200413f237f41ab11ad6d161bc7239c84dcb631773ccd7de3dfe4b5c267c" +checksum = "da0df0e5185db44f69b44f26786fe401b6c293d1907744beaa7fa62b2e5a517a" dependencies = [ "autocfg", ] @@ -2676,6 +2742,12 @@ dependencies = [ "windows-targets 0.48.5", ] +[[package]] +name = "paste" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de3145af08024dea9fa9914f381a17b8fc6034dfb00f3a84013f7ff43f29ed4c" + [[package]] name = "pem" version = "3.0.3" @@ -3971,6 +4043,37 @@ dependencies = [ "threadpool", ] +[[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.31" diff --git a/dragonfly-client/Cargo.toml b/dragonfly-client/Cargo.toml index 334c0dc1..bc23ab4e 100644 --- a/dragonfly-client/Cargo.toml +++ b/dragonfly-client/Cargo.toml @@ -78,3 +78,8 @@ tokio-rustls = "0.25.0-alpha.4" http-body-util = "0.1.0" futures-util = "0.3.30" termion = "3.0.0" +[target.'cfg(not(target_env = "msvc"))'.dependencies] +tikv-jemallocator = { version = "0.5.4", features = ["profiling", "unprefixed_malloc_on_supported_platforms"] } + +[target.'cfg(target_os = "linux")'.dependencies] +jemalloc_pprof = "0.1.0" diff --git a/dragonfly-client/src/bin/dfdaemon/main.rs b/dragonfly-client/src/bin/dfdaemon/main.rs index fdf30d95..58247be4 100644 --- a/dragonfly-client/src/bin/dfdaemon/main.rs +++ b/dragonfly-client/src/bin/dfdaemon/main.rs @@ -39,6 +39,14 @@ use std::sync::Arc; use tokio::sync::mpsc; use tracing::{error, info, Level}; +#[cfg(not(target_env = "msvc"))] +#[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:true,lg_prof_sample:19\0"; + #[derive(Debug, Parser)] #[command( name = dfdaemon::NAME, diff --git a/dragonfly-client/src/stats/mod.rs b/dragonfly-client/src/stats/mod.rs index c849c372..a6e276c2 100644 --- a/dragonfly-client/src/stats/mod.rs +++ b/dragonfly-client/src/stats/mod.rs @@ -30,10 +30,10 @@ const DEFAULT_PROFILER_SECONDS: u64 = 10; // DEFAULT_PROFILER_FREQUENCY is the default frequency to start profiling. const DEFAULT_PROFILER_FREQUENCY: i32 = 1000; -// PProfQueryParams is the query params to start profiling. +// PProfProfileQueryParams is the query params to start profiling. #[derive(Deserialize, Serialize)] #[serde(default)] -pub struct PProfQueryParams { +pub struct PProfProfileQueryParams { // seconds is the seconds to start profiling. pub seconds: u64, @@ -41,8 +41,8 @@ pub struct PProfQueryParams { pub frequency: i32, } -// PProfQueryParams implements the default. -impl Default for PProfQueryParams { +// PProfProfileQueryParams implements the default. +impl Default for PProfProfileQueryParams { fn default() -> Self { Self { seconds: DEFAULT_PROFILER_SECONDS, @@ -84,16 +84,24 @@ impl Stats { // Clone the shutdown channel. let mut shutdown = self.shutdown.clone(); - // Create the stats route. - let stats_route = warp::path!("debug" / "pprof" / "profile") + // Create the pprof profile route. + let pprof_profile_route = warp::path!("debug" / "pprof" / "profile") .and(warp::get()) - .and(warp::query::()) - .and_then(Self::stats_handler); + .and(warp::query::()) + .and_then(Self::pprof_profile_handler); + + // Create the pprof heap route. + let pprof_heap_route = warp::path!("debug" / "pprof" / "heap") + .and(warp::get()) + .and_then(Self::pprof_heap_handler); + + // Create the pprof routes. + let pprof_routes = pprof_profile_route.or(pprof_heap_route); // Start the stats server and wait for it to finish. info!("stats server listening on {}", self.addr); tokio::select! { - _ = warp::serve(stats_route).run(self.addr) => { + _ = warp::serve(pprof_routes).run(self.addr) => { // Stats server ended. info!("stats server ended"); } @@ -105,35 +113,29 @@ impl Stats { } // stats_handler handles the stats request. - async fn stats_handler(query_params: PProfQueryParams) -> Result { + async fn pprof_profile_handler( + query_params: PProfProfileQueryParams, + ) -> Result { info!( "start profiling for {} seconds with {} frequency", query_params.seconds, query_params.frequency ); - let guard = match ProfilerGuard::new(query_params.frequency) { - Ok(guard) => guard, - Err(err) => { - error!("failed to create profiler guard: {}", err); - return Err(warp::reject::reject()); - } - }; + + let guard = ProfilerGuard::new(query_params.frequency).map_err(|err| { + error!("failed to create profiler guard: {}", err); + warp::reject::reject() + })?; tokio::time::sleep(Duration::from_secs(query_params.seconds)).await; - let report = match guard.report().build() { - Ok(report) => report, - Err(err) => { - error!("failed to build profiler report: {}", err); - return Err(warp::reject::reject()); - } - }; + let report = guard.report().build().map_err(|err| { + error!("failed to build profiler report: {}", err); + warp::reject::reject() + })?; - let profile = match report.pprof() { - Ok(profile) => profile, - Err(err) => { - error!("failed to get pprof profile: {}", err); - return Err(warp::reject::reject()); - } - }; + let profile = report.pprof().map_err(|err| { + error!("failed to get pprof profile: {}", err); + warp::reject::reject() + })?; let mut body: Vec = Vec::new(); profile.write_to_vec(&mut body).map_err(|err| { @@ -143,4 +145,26 @@ impl Stats { Ok(body) } + + // pprof_heap_handler handles the pprof heap request. + async fn pprof_heap_handler() -> Result { + info!("start heap profiling"); + #[cfg(target_os = "linux")] + { + let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await; + if !prof_ctl.activated() { + return Err(warp::reject::reject()); + } + + let pprof = prof_ctl.dump_pprof().map_err(|err| { + error!("failed to dump pprof: {}", err); + warp::reject::reject() + })?; + + Ok(pprof) + } + + #[cfg(not(target_os = "linux"))] + Err::(warp::reject::reject()) + } }