Add flamegraph feature

Author: erikgrinaker

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]

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<impl IntoResponse, (StatusCode, String)> {
+    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.

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"]

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<impl IntoResponse, (StatusCode, String)
     Ok(pprof)
 }
 
+#[cfg(feature = "flamegraph")]
+pub async fn handle_get_heap_flamegraph() -> Result<impl IntoResponse, (StatusCode, String)> {
+    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,

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<Vec<u8>> {
+        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<Vec<u8>> {
+        let f = self.dump()?;
+        let dump_reader = BufReader::new(f);
+        let profile = parse_jeheap(dump_reader, MAPPINGS.as_deref())?;
+        profile.to_flamegraph(opts)
+    }
 }

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"]

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<String>,
     ) -> Vec<u8> {
-        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<String>,
+    ) -> 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<Vec<u8>> {
+        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<u64, proto::Location> =
+            profile.location.into_iter().map(|l| (l.id, l)).collect();
+        let functions: HashMap<u64, proto::Function> =
+            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<Vec<&str>, 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::<Vec<_>>();
+        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)
     }
 }