From 5945bbeab1d167418d700cfe2d918421af5e702f Mon Sep 17 00:00:00 2001
From: Roman Gershman <romange@gmail.com>
Date: Mon, 18 Oct 2021 11:13:23 +0300
Subject: [PATCH] Allow skipping stack frames after a specified frame name

Fixes #221
---
 src/collapse/perf.rs | 62 +++++++++++++++++++++++++++++++++++++-------
 1 file changed, 53 insertions(+), 9 deletions(-)

diff --git a/src/collapse/perf.rs b/src/collapse/perf.rs
index 16fe6f52..8d6adbfa 100644
--- a/src/collapse/perf.rs
+++ b/src/collapse/perf.rs
@@ -23,6 +23,13 @@ mod logging {
     }
 }
 
+#[derive(PartialEq)]
+enum SkipStackOpt {
+    DontSkip,
+    Entirely,
+    Partially,
+}
+
 /// `perf` folder configuration options.
 #[derive(Clone, Debug)]
 #[non_exhaustive]
@@ -63,6 +70,11 @@ pub struct Options {
     ///
     /// Default is the number of logical cores on your machine.
     pub nthreads: usize,
+
+    /// Skips stack lines after the specified string is matched as substring of a function name.
+    /// In case no function is matched the whole stack is returned.
+    /// Default is `None`.
+    pub skip_after: Option<String>,
 }
 
 impl Default for Options {
@@ -75,6 +87,7 @@ impl Default for Options {
             include_pid: false,
             include_tid: false,
             nthreads: *common::DEFAULT_NTHREADS,
+            skip_after: None,
         }
     }
 }
@@ -110,8 +123,8 @@ pub struct Folder {
     /// Called pname after original stackcollapse-perf source.
     pname: String,
 
-    /// Skip all stack lines in this event.
-    skip_stack: bool,
+    /// Whether to skip stack lines in this event.
+    skip_stack: SkipStackOpt,
 
     /// Function entries on the stack in this entry thus far.
     stack: VecDeque<String>,
@@ -132,7 +145,7 @@ impl From<Options> for Folder {
             in_event: false,
             nstacks_per_job: common::DEFAULT_NSTACKS_PER_JOB,
             pname: String::default(),
-            skip_stack: false,
+            skip_stack: SkipStackOpt::DontSkip,
             stack: VecDeque::default(),
             opt,
         }
@@ -189,7 +202,7 @@ impl CollapsePrivate for Folder {
 
         // Reset state...
         self.in_event = false;
-        self.skip_stack = false;
+        self.skip_stack = SkipStackOpt::DontSkip;
         self.stack.clear();
         Ok(())
     }
@@ -246,7 +259,7 @@ impl CollapsePrivate for Folder {
             in_event: false,
             nstacks_per_job: self.nstacks_per_job,
             pname: String::new(),
-            skip_stack: false,
+            skip_stack: SkipStackOpt::DontSkip,
             stack: VecDeque::default(),
             opt: self.opt.clone(),
         }
@@ -367,7 +380,7 @@ impl Folder {
             if let Some(event) = event {
                 if let Some(ref event_filter) = self.event_filter {
                     if event != event_filter {
-                        self.skip_stack = true;
+                        self.skip_stack = SkipStackOpt::Entirely;
                         return;
                     }
                 } else {
@@ -465,7 +478,7 @@ impl Folder {
     //     7f53389994d0 [unknown] ([unknown])
     //                0 [unknown] ([unknown])
     fn on_stack_line(&mut self, line: &str) {
-        if self.skip_stack {
+        if self.skip_stack == SkipStackOpt::Entirely || self.skip_stack == SkipStackOpt::Partially {
             return;
         }
 
@@ -531,6 +544,12 @@ impl Folder {
             while let Some(func) = self.cache_line.pop() {
                 self.stack.push_front(func);
             }
+
+            if let Some(skip_after) = &self.opt.skip_after {
+                if rawfunc.eq(skip_after) {
+                    self.skip_stack = SkipStackOpt::Partially;
+                }
+            }
         } else {
             logging::weird_stack_line(line);
         }
@@ -538,7 +557,7 @@ impl Folder {
 
     fn after_event(&mut self, occurrences: &mut Occurrences) {
         // end of stack, so emit stack entry
-        if !self.skip_stack {
+        if self.skip_stack == SkipStackOpt::Partially || self.skip_stack == SkipStackOpt::DontSkip {
             // allocate a string that is long enough to hold the entire stack string
             let mut stack_str = String::with_capacity(
                 self.pname.len() + self.stack.iter().fold(0, |a, s| a + s.len() + 1),
@@ -558,7 +577,7 @@ impl Folder {
 
         // reset for the next event
         self.in_event = false;
-        self.skip_stack = false;
+        self.skip_stack = SkipStackOpt::DontSkip;
         self.stack.clear();
     }
 }
@@ -755,6 +774,30 @@ mod tests {
         <Folder as Collapse>::collapse(&mut folder, &bytes[..], io::sink())
     }
 
+    #[test]
+    fn test_skip_after() -> io::Result<()> {
+        let path = "./tests/data/collapse-perf/go-stacks.txt";
+        let mut file = fs::File::open(path)?;
+        let mut bytes = Vec::new();
+        file.read_to_end(&mut bytes)?;
+        let mut folder = {
+            let options = Options {
+                skip_after: Some("main.init".to_string()),
+                ..Default::default()
+            };
+            Folder::from(options)
+        };
+        let mut buf_actual = Vec::new();
+        <Folder as Collapse>::collapse(&mut folder, &bytes[..], &mut buf_actual)?;
+        let lines = std::str::from_utf8(&buf_actual[..]).unwrap().lines();
+        for line in lines {
+            if line.contains("main.init") {
+                assert!(line.contains("go;main.init;")); // we removed the frames above "main.init"
+            }
+        }
+        Ok(())
+    }
+
     /// Varies the nstacks_per_job parameter and outputs the 10 fastests configurations by file.
     ///
     /// Command: `cargo test bench_nstacks_perf --release -- --ignored --nocapture`
@@ -791,6 +834,7 @@ mod tests {
                 include_pid: rng.gen(),
                 include_tid: rng.gen(),
                 nthreads: rng.gen_range(2..=32),
+                skip_after: None,
             };
 
             for (path, input) in inputs.iter() {