Skip to content

Add flamegraph feature #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 3, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
@@ -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]
37 changes: 37 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
@@ -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.
3 changes: 3 additions & 0 deletions example/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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"]
24 changes: 24 additions & 0 deletions example/src/main.rs
Original file line number Diff line number Diff line change
@@ -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,
23 changes: 23 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}
2 changes: 2 additions & 0 deletions util/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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"]
73 changes: 66 additions & 7 deletions util/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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)
}
}