From ea5fc992be98b56b40a7ee4271b9dac857fd63a0 Mon Sep 17 00:00:00 2001 From: Nils Diedrich Date: Tue, 5 May 2026 16:41:23 +0200 Subject: [PATCH] feat(profiling): add --show-from filter to profiling aggregate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Datadog Continuous Profiler UI exposes a `show_from()` flame-graph filter that zooms the displayed graph to subtrees rooted at a chosen function. The filter is applied client-side — the API still returns the full aggregated profile — so reproducing it from the CLI required hand-rolled jq, which is awkward for AI agents and humans alike. This adds an equivalent native flag. `pup profiling aggregate --show-from=` does, in order: 1. Resolves the function name to string IDs (exact match against `data.strings`). 2. Walks the response's frame table to find every `frame_id` whose `function` field (`frameSchema[4]`) refers to one of those string IDs. A single logical function commonly maps to many frame IDs after Go inlining / generics. 3. Walks `flameGraph` and collects the topmost subtree rooted at any matching frame on each call path (does not descend into nested re-entries — matches UI semantics). 4. Merges the collected subtrees into a single root by recursively collapsing children that share the same function name. Without this, the same logical function appears as N siblings at depth 0 (one per inlined frame ID), which doesn't match the UI. 5. Prunes `frames` and `strings` to only the entries the trimmed subtree references (with consistent ID remapping), and drops UI aggregation fields (`metadata`, `summaryTable`, `endpointCounts`, etc.) that describe the unfiltered profile rather than the filtered subtree. Concrete impact on a 14-hour `service:logs-event-store-reader` query: unfiltered response is 43 MB; with `--show-from=vectorEQStringNotMulti` the response drops to ~107 KB and shows the same call tree as the UI. - New `--show-from` flag on `ProfilingActions::Aggregate` (src/main.rs) - `apply_show_from`, `merge_subtrees_by_function`, and `prune_aggregate_response` helpers in src/commands/profiling.rs Co-Authored-By: Claude Opus 4.7 (1M context) --- src/commands/profiling.rs | 287 +++++++++++++++++++++++++++++++++++++- src/main.rs | 7 + 2 files changed, 293 insertions(+), 1 deletion(-) diff --git a/src/commands/profiling.rs b/src/commands/profiling.rs index 3c454390..cfb0ffe2 100644 --- a/src/commands/profiling.rs +++ b/src/commands/profiling.rs @@ -57,6 +57,7 @@ pub async fn aggregate( to: String, limit: u32, aggregation_function: String, + show_from: Option, ) -> Result<()> { let (from_iso, to_iso) = parse_window(&from, &to)?; // /profiling/api/v1/aggregate expects a flat body — query/from/to are siblings, not wrapped in filter{}. @@ -68,12 +69,296 @@ pub async fn aggregate( "limit": limit, "aggregationFunction": aggregation_function, }); - let resp = client::raw_post(cfg, "/profiling/api/v1/aggregate", body) + let mut resp = client::raw_post(cfg, "/profiling/api/v1/aggregate", body) .await .map_err(|e| anyhow::anyhow!("failed to aggregate profiles: {e:?}"))?; + if let Some(name) = show_from { + apply_show_from(&mut resp, &name)?; + prune_aggregate_response(&mut resp); + } formatter::output(cfg, &resp) } +// When --show-from has trimmed `flameGraph` to a small subtree, the rest of the +// response (frames, strings, metadata) is mostly dead weight. Compact `frames` +// and `strings` to just the entries the trimmed flame graph still references, +// remap IDs accordingly, and drop heavy fields the UI uses but a CLI consumer +// almost never wants alongside a single-function view. +fn prune_aggregate_response(resp: &mut serde_json::Value) { + use std::collections::{HashMap, HashSet}; + + let Some(flame) = resp.get("flameGraph").cloned() else { + return; + }; + let Some(frames) = resp.get("frames").and_then(|f| f.as_array()).cloned() else { + return; + }; + let Some(strings) = resp.get("strings").and_then(|s| s.as_array()).cloned() else { + return; + }; + + let mut used_frames: HashSet = HashSet::new(); + collect_used_frame_ids(&flame, &mut used_frames); + if used_frames.is_empty() { + return; + } + + // Indices 2..=5 of a frame are string-table refs (library, package, function, file). + let mut used_strings: HashSet = HashSet::new(); + for &fid in &used_frames { + let Some(f) = frames.get(fid as usize).and_then(|f| f.as_array()) else { + continue; + }; + for idx in 2usize..=5 { + if let Some(sid) = f.get(idx).and_then(|v| v.as_i64()) { + used_strings.insert(sid); + } + } + } + + let mut sorted_frames: Vec = used_frames.iter().copied().collect(); + sorted_frames.sort_unstable(); + let frame_remap: HashMap = sorted_frames + .iter() + .enumerate() + .map(|(i, &old)| (old, i as i64)) + .collect(); + + let mut sorted_strings: Vec = used_strings.iter().copied().collect(); + sorted_strings.sort_unstable(); + let string_remap: HashMap = sorted_strings + .iter() + .enumerate() + .map(|(i, &old)| (old, i as i64)) + .collect(); + + let new_frames: Vec = sorted_frames + .iter() + .filter_map(|&old_fid| { + let mut f = frames.get(old_fid as usize).cloned()?; + if let Some(arr) = f.as_array_mut() { + for idx in 2usize..=5 { + if let Some(slot) = arr.get_mut(idx) { + if let Some(old_sid) = slot.as_i64() { + if let Some(&new_sid) = string_remap.get(&old_sid) { + *slot = json!(new_sid); + } + } + } + } + } + Some(f) + }) + .collect(); + + let new_strings: Vec = sorted_strings + .iter() + .filter_map(|&old_sid| strings.get(old_sid as usize).cloned()) + .collect(); + + let mut new_flame = flame; + remap_frame_ids(&mut new_flame, &frame_remap); + + resp["flameGraph"] = new_flame; + resp["frames"] = json!(new_frames); + resp["strings"] = json!(new_strings); + + // These are large UI-oriented fields with no relationship to the + // post-filter subtree; strip to keep terminal output usable. + if let Some(obj) = resp.as_object_mut() { + for k in [ + "metadata", + "frameSchemas", + "endpointCounts", + "endpointValues", + "summaryTable", + "summaryValues", + "summaryDurations", + "availableAttributes", + "featureUpgrades", + "languageFrameCounts", + "profileIds", + "emptyStateReason", + ] { + obj.remove(k); + } + } +} + +fn collect_used_frame_ids(node: &serde_json::Value, out: &mut std::collections::HashSet) { + let Some(arr) = node.as_array() else { + return; + }; + if let Some(fid) = arr.first().and_then(|v| v.as_i64()) { + out.insert(fid); + } + if let Some(children) = arr.get(3).and_then(|c| c.as_array()) { + for c in children { + collect_used_frame_ids(c, out); + } + } +} + +fn remap_frame_ids(node: &mut serde_json::Value, remap: &std::collections::HashMap) { + let Some(arr) = node.as_array_mut() else { + return; + }; + if let Some(slot) = arr.first_mut() { + if let Some(old) = slot.as_i64() { + if let Some(&new) = remap.get(&old) { + *slot = json!(new); + } + } + } + if let Some(children) = arr.get_mut(3).and_then(|c| c.as_array_mut()) { + for c in children { + remap_frame_ids(c, remap); + } + } +} + +// Client-side equivalent of the UI's `show_from()` flame-graph filter. +// Replaces `data.flameGraph` with a synthetic root whose children are the +// topmost subtrees rooted at frames with the given function name (exact match). +// Frame names are resolved via `data.frames[i][4]` (the `function` field of +// `frameSchema`) → `data.strings[id]`. `data.frames` and `data.strings` are +// left untouched so callers can still resolve frame metadata. +fn apply_show_from(resp: &mut serde_json::Value, function_name: &str) -> Result<()> { + use std::collections::HashSet; + + let strings = resp + .get("strings") + .and_then(|s| s.as_array()) + .ok_or_else(|| anyhow::anyhow!("response has no 'strings' array"))?; + let target_string_ids: HashSet = strings + .iter() + .enumerate() + .filter_map(|(i, s)| { + s.as_str() + .filter(|st| *st == function_name) + .map(|_| i as i64) + }) + .collect(); + if target_string_ids.is_empty() { + anyhow::bail!( + "--show-from: no string in 'strings' matches function name {function_name:?}" + ); + } + + let frames_arr: Vec = resp + .get("frames") + .and_then(|f| f.as_array()) + .ok_or_else(|| anyhow::anyhow!("response has no 'frames' array"))? + .clone(); + let target_frame_ids: HashSet = frames_arr + .iter() + .enumerate() + .filter_map(|(i, f)| { + let arr = f.as_array()?; + let fname_id = arr.get(4)?.as_i64()?; + if target_string_ids.contains(&fname_id) { + Some(i as i64) + } else { + None + } + }) + .collect(); + if target_frame_ids.is_empty() { + anyhow::bail!( + "--show-from: no frame in 'frames' references function name {function_name:?}" + ); + } + + let flame = resp + .get("flameGraph") + .cloned() + .ok_or_else(|| anyhow::anyhow!("response has no 'flameGraph'"))?; + let mut matches = Vec::new(); + collect_show_from_subtrees(&flame, &target_frame_ids, &mut matches); + if matches.is_empty() { + anyhow::bail!( + "--show-from: no node in 'flameGraph' references function name {function_name:?}" + ); + } + + resp["flameGraph"] = merge_subtrees_by_function(matches, &frames_arr); + Ok(()) +} + +// Merge a set of flame-graph subtrees into a single subtree by collapsing +// children that share the same function name (frames[fid][4] -> string id). +// Datadog's flame graph emits a distinct `frame_id` per (function, file, line) +// tuple, so a single logical function can appear under many frame IDs after +// inlining/generics. The UI merges siblings by display name; we mirror that. +fn merge_subtrees_by_function( + mut nodes: Vec, + frames: &[serde_json::Value], +) -> serde_json::Value { + if nodes.len() == 1 { + return nodes.pop().unwrap(); + } + let representative_fid = nodes + .iter() + .filter_map(|n| n.as_array().and_then(|a| a.first()).and_then(|v| v.as_i64())) + .next() + .unwrap_or(0); + let total_value: i64 = nodes + .iter() + .filter_map(|n| n.as_array().and_then(|a| a.get(1)).and_then(|v| v.as_i64())) + .sum(); + let mut all_children: Vec = Vec::new(); + for n in &nodes { + if let Some(children) = n.as_array().and_then(|a| a.get(3)).and_then(|c| c.as_array()) { + all_children.extend(children.iter().cloned()); + } + } + // Group by function string id; preserve insertion order so output is stable. + use std::collections::BTreeMap; + let mut groups: BTreeMap> = BTreeMap::new(); + let mut keyless: Vec = Vec::new(); + for c in all_children { + let key = c + .as_array() + .and_then(|a| a.first()) + .and_then(|v| v.as_i64()) + .and_then(|fid| frames.get(fid as usize)) + .and_then(|f| f.as_array()) + .and_then(|fa| fa.get(4)) + .and_then(|v| v.as_i64()); + match key { + Some(k) => groups.entry(k).or_default().push(c), + None => keyless.push(c), + } + } + let mut merged_children: Vec = groups + .into_values() + .map(|grp| merge_subtrees_by_function(grp, frames)) + .collect(); + merged_children.append(&mut keyless); + json!([representative_fid, total_value, -1.0, merged_children]) +} + +fn collect_show_from_subtrees( + node: &serde_json::Value, + targets: &std::collections::HashSet, + out: &mut Vec, +) { + let Some(arr) = node.as_array() else { + return; + }; + if let Some(fid) = arr.first().and_then(|v| v.as_i64()) { + if targets.contains(&fid) { + out.push(node.clone()); + return; // topmost match only — don't descend into nested re-entries + } + } + if let Some(children) = arr.get(3).and_then(|c| c.as_array()) { + for c in children { + collect_show_from_subtrees(c, targets, out); + } + } +} + pub async fn analysis(cfg: &Config, profile_id: &str, event_id: Option) -> Result<()> { let path = format!("/profiling/api/v1/profiles/{profile_id}/analysis"); let query: Vec<(&str, &str)> = match event_id.as_deref() { diff --git a/src/main.rs b/src/main.rs index d5ed990f..ffede9cb 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8868,6 +8868,11 @@ enum ProfilingActions { limit: u32, #[arg(long, default_value = "sum", help = "Aggregation function: sum or avg")] aggregation_function: String, + #[arg( + long, + help = "Filter the flame graph to subtrees rooted at frames whose function name matches exactly (UI equivalent: show_from())" + )] + show_from: Option, }, /// Get automated analysis for a profile Analysis { @@ -13984,6 +13989,7 @@ async fn main_inner() -> anyhow::Result<()> { to, limit, aggregation_function, + show_from, } => { commands::profiling::aggregate( &cfg, @@ -13993,6 +13999,7 @@ async fn main_inner() -> anyhow::Result<()> { to, limit, aggregation_function, + show_from, ) .await?; }