diff --git a/Cargo.toml b/Cargo.toml index 6e669b7..fa6b831 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,9 @@ categories = [ documentation = "https://docs.rs/jemalloc_pprof/latest/jemalloc_pprof/" homepage = "https://crates.io/crates/jemalloc_pprof" +[package.metadata.docs.rs] +all-features = true + [workspace.dependencies] anyhow = "1" flate2 = "1.0" @@ -39,6 +42,7 @@ errno = "0.3" util = { path = "./util", version = "0.6", package = "pprof_util" } mappings = { path = "./mappings", version = "0.6" } backtrace = "0.3" +inferno = "0.12" [dependencies] util.workspace = true @@ -52,6 +56,7 @@ tempfile.workspace = true tokio.workspace = true [features] +flamegraph = ["util/flamegraph"] symbolize = ["util/symbolize"] [dev-dependencies] diff --git a/README.md b/README.md index f75ea4a..29b23f1 100644 --- a/README.md +++ b/README.md @@ -102,6 +102,43 @@ To generate symbolized profiles, enable the `symbolize` crate feature: jemalloc_pprof = { version = "0.6", features = ["symbolize"] } ``` +### Flamegraph SVGs + +The `flamegraph` crate feature can also be enabled to generate interactive flamegraph SVGs directly +(implies the `symbolize` feature): + +```toml +jemalloc_pprof = { version = "0.6", features = ["flamegraph"] } +``` + +We can then adjust the example above to also emit a flamegraph SVG: + +```rust,ignore +#[tokio::main] +async fn main() { + let app = axum::Router::new() + .route("/debug/pprof/heap", axum::routing::get(handle_get_heap)) + .route("/debug/pprof/heap/flamegraph", axum::routing::get(handle_get_heap_flamegraph)); + // ... +} + +pub async fn handle_get_heap_flamegraph() -> Result { + use axum::body::Body; + use axum::http::header::CONTENT_TYPE; + use axum::response::Response; + + let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await; + require_profiling_activated(&prof_ctl)?; + let svg = prof_ctl + .dump_flamegraph() + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + Response::builder() + .header(CONTENT_TYPE, "image/svg+xml") + .body(Body::from(svg)) + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string())) +} +``` + ### Writeable temporary directory The way this library works is that it creates a new temporary file (in the [platform-specific default temp dir](https://docs.rs/tempfile/latest/tempfile/struct.NamedTempFile.html)), and instructs jemalloc to dump a profile into that file. Therefore the platform respective temporary directory must be writeable by the process. After reading and converting it to pprof, the file is cleaned up via the destructor. A single profile tends to be only a few kilobytes large, so it doesn't require a significant space, but it's non-zero and needs to be writeable. diff --git a/example/Cargo.toml b/example/Cargo.toml index fbb10f1..ca925b9 100644 --- a/example/Cargo.toml +++ b/example/Cargo.toml @@ -10,3 +10,6 @@ tokio = { version = "1", features = ["full"] } axum = "0.7.2" [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemallocator = { version = "0.6", features = ["profiling", "stats", "unprefixed_malloc_on_supported_platforms", "background_threads"] } + +[features] +flamegraph = ["jemalloc_pprof/flamegraph"] \ No newline at end of file diff --git a/example/src/main.rs b/example/src/main.rs index eed32d1..7d5e604 100644 --- a/example/src/main.rs +++ b/example/src/main.rs @@ -18,6 +18,13 @@ async fn main() { let app = axum::Router::new().route("/debug/pprof/heap", axum::routing::get(handle_get_heap)); + // Add a flamegraph SVG route if enabled via `cargo run -F flamegraph`. + #[cfg(feature = "flamegraph")] + let app = app.route( + "/debug/pprof/heap/flamegraph", + axum::routing::get(handle_get_heap_flamegraph), + ); + // run our app with hyper, listening globally on port 3000 let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap(); axum::serve(listener, app).await.unwrap(); @@ -32,6 +39,23 @@ pub async fn handle_get_heap() -> Result Result { + use axum::body::Body; + use axum::http::header::CONTENT_TYPE; + use axum::response::Response; + + let mut prof_ctl = jemalloc_pprof::PROF_CTL.as_ref().unwrap().lock().await; + require_profiling_activated(&prof_ctl)?; + let svg = prof_ctl + .dump_flamegraph() + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string()))?; + Response::builder() + .header(CONTENT_TYPE, "image/svg+xml") + .body(Body::from(svg)) + .map_err(|err| (StatusCode::INTERNAL_SERVER_ERROR, err.to_string())) +} + /// Checks whether jemalloc profiling is activated an returns an error response if not. fn require_profiling_activated( prof_ctl: &jemalloc_pprof::JemallocProfCtl, diff --git a/src/lib.rs b/src/lib.rs index 468f759..a758b29 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -29,6 +29,8 @@ use tempfile::NamedTempFile; use tikv_jemalloc_ctl::raw; use tokio::sync::Mutex; +#[cfg(feature = "flamegraph")] +pub use util::FlamegraphOptions; use util::{parse_jeheap, ProfStartTime}; /// Activate jemalloc profiling. @@ -166,4 +168,25 @@ impl JemallocProfCtl { let pprof = profile.to_pprof(("inuse_space", "bytes"), ("space", "bytes"), None); Ok(pprof) } + + /// Dump a profile flamegraph in SVG format. + #[cfg(feature = "flamegraph")] + pub fn dump_flamegraph(&mut self) -> anyhow::Result> { + let mut opts = FlamegraphOptions::default(); + opts.title = "inuse_space".to_string(); + opts.count_name = "bytes".to_string(); + self.dump_flamegraph_with_options(&mut opts) + } + + /// Dump a profile flamegraph in SVG format with the given options. + #[cfg(feature = "flamegraph")] + pub fn dump_flamegraph_with_options( + &mut self, + opts: &mut FlamegraphOptions, + ) -> anyhow::Result> { + let f = self.dump()?; + let dump_reader = BufReader::new(f); + let profile = parse_jeheap(dump_reader, MAPPINGS.as_deref())?; + profile.to_flamegraph(opts) + } } diff --git a/util/Cargo.toml b/util/Cargo.toml index 8fba4ef..2c185fd 100644 --- a/util/Cargo.toml +++ b/util/Cargo.toml @@ -17,6 +17,8 @@ anyhow.workspace = true num.workspace = true paste.workspace = true backtrace = { workspace = true, optional = true } +inferno = { workspace = true, optional = true } [features] +flamegraph = ["symbolize", "dep:inferno"] symbolize = ["dep:backtrace"] diff --git a/util/src/lib.rs b/util/src/lib.rs index d519c2f..a69fa46 100644 --- a/util/src/lib.rs +++ b/util/src/lib.rs @@ -15,6 +15,9 @@ use prost::Message; pub use cast::CastFrom; pub use cast::TryCastFrom; +#[cfg(feature = "flamegraph")] +pub use inferno::flamegraph::Options as FlamegraphOptions; + /// Start times of the profiler. #[derive(Copy, Clone, Debug)] pub enum ProfStartTime { @@ -51,7 +54,7 @@ impl StringTable { } #[path = "perftools.profiles.rs"] -mod pprof_types; +mod proto; /// A single sample in the profile. The stack is a list of addresses. #[derive(Clone, Debug)] @@ -104,8 +107,21 @@ impl StackProfile { period_type: (&str, &str), anno_key: Option, ) -> Vec { - use crate::pprof_types as proto; + let profile = self.to_pprof_proto(sample_type, period_type, anno_key); + let encoded = profile.encode_to_vec(); + + let mut gz = GzEncoder::new(Vec::new(), Compression::default()); + gz.write_all(&encoded).unwrap(); + gz.finish().unwrap() + } + /// Converts the profile into the pprof Protobuf format (see `pprof/profile.proto`). + fn to_pprof_proto( + &self, + sample_type: (&str, &str), + period_type: (&str, &str), + anno_key: Option, + ) -> proto::Profile { let mut profile = proto::Profile::default(); let mut strings = StringTable::new(); @@ -192,7 +208,7 @@ impl StackProfile { let addr = u64::cast_from(*addr) - 1; let loc_id = *location_ids.entry(addr).or_insert_with(|| { - // pprof_types.proto says the location id may be the address, but Polar Signals + // profile.proto says the location id may be the address, but Polar Signals // insists that location ids are sequential, starting with 1. let id = u64::cast_from(profile.location.len()) + 1; @@ -275,11 +291,54 @@ impl StackProfile { profile.string_table = strings.finish(); - let encoded = profile.encode_to_vec(); + profile + } - let mut gz = GzEncoder::new(Vec::new(), Compression::default()); - gz.write_all(&encoded).unwrap(); - gz.finish().unwrap() + /// Converts the profile into a flamegraph SVG, using the given options. + #[cfg(feature = "flamegraph")] + pub fn to_flamegraph(&self, opts: &mut FlamegraphOptions) -> anyhow::Result> { + use std::collections::HashMap; + + // We start from a symbolized Protobuf profile. We just pass in empty type names, since + // they're not used in the final flamegraph. + let profile = self.to_pprof_proto(("", ""), ("", ""), None); + + // Index locations, functions, and strings. + let locations: HashMap = + profile.location.into_iter().map(|l| (l.id, l)).collect(); + let functions: HashMap = + profile.function.into_iter().map(|f| (f.id, f)).collect(); + let strings = profile.string_table; + + // Resolve stacks as function name vectors, and sum sample values per stack. Also reverse + // the stack, since inferno expects it bottom-up. + let mut stacks: HashMap, i64> = HashMap::new(); + for sample in profile.sample { + let mut stack = Vec::with_capacity(sample.location_id.len()); + for location in sample.location_id.into_iter().rev() { + let location = locations.get(&location).expect("missing location"); + for line in location.line.iter().rev() { + let function = functions.get(&line.function_id).expect("missing function"); + let name = strings.get(function.name as usize).expect("missing string"); + stack.push(name.as_str()); + } + } + let value = sample.value.first().expect("missing value"); + *stacks.entry(stack).or_default() += value; + } + + // Construct stack lines for inferno. + let mut lines = stacks + .into_iter() + .map(|(stack, value)| format!("{} {}", stack.join(";"), value)) + .collect::>(); + lines.sort(); + + // Generate the flamegraph SVG. + let mut bytes = Vec::new(); + let lines = lines.iter().map(|line| line.as_str()); + inferno::flamegraph::from_lines(opts, lines, &mut bytes)?; + Ok(bytes) } }