Add an option to filter and truncate stacks to a provided symbol

src/bin/flamegraph.rs

@@ -212,6 +212,13 @@ struct Opt {
     #[clap(long = "width", value_name = "UINT")]
     width: Option<usize>,
 
+    /// Omit samples whose stacks do not contain this symbol. When this symbol is in a sample's
+    /// stack, truncate the call stack so that this is the bottom-most symbol.
+    /// This is particularly useful when you want to profile a specific function in a codebase that
+    /// uses some kind of heavy runtime like rayon, tokio, or the rustc query system.
+    #[clap(long = "base", value_name = "STRING")]
+    base: Vec<String>,
+
     // ************ //
     // *** ARGS *** //
     // ************ //
@@ -256,6 +263,7 @@ impl<'a> Opt {
         options.color_diffusion = self.color_diffusion;
         options.reverse_stack_order = self.reverse;
         options.flame_chart = self.flame_chart;
+        options.base = self.base;
 
         if self.flame_chart && self.title == defaults::TITLE {
             options.title = defaults::CHART_TITLE.to_owned();

src/flamegraph/mod.rs

@@ -258,6 +258,9 @@ pub struct Options<'a> {
     ///
     /// Note that stack is not sorted and will be reversed
     pub flame_chart: bool,
+
+    /// Base symbols
+    pub base: Vec<String>,
 }
 
 impl<'a> Options<'a> {
@@ -322,6 +325,7 @@ impl<'a> Default for Options<'a> {
             no_javascript: Default::default(),
             color_diffusion: Default::default(),
             flame_chart: Default::default(),
+            base: Default::default(),
 
             #[cfg(feature = "nameattr")]
             func_frameattrs: Default::default(),
@@ -450,7 +454,28 @@ where
         merge::frames(lines, false)?
     } else {
         // Sort lines by default.
-        let mut lines: Vec<&str> = lines.into_iter().collect();
+        let mut lines: Vec<&str> = if opt.base.is_empty() {
+            lines.into_iter().collect()
+        } else {
+            lines
+                .into_iter()
+                .filter_map(|line| {
+                    let mut cursor = line.len();
+                    for symbol in line.rsplit(';') {
+                        cursor -= symbol.len();
+                        if opt.base.iter().any(|b| b == symbol) {
+                            break;
+                        }
+                        cursor = cursor.saturating_sub(1);
+                    }
+                    if cursor == 0 {
+                        None
+                    } else {
+                        Some(&line[cursor..])
+                    }
+                })
+                .collect()
+        };
         lines.sort_unstable();
         merge::frames(lines, false)?
     };