Port stackcollapse-recursive.pl

Cargo.toml

@@ -102,6 +102,11 @@ name = "inferno-collapse-guess"
 path = "src/bin/collapse-guess.rs"
 required-features = ["cli"]
 
+[[bin]]
+name = "inferno-collapse-recursive"
+path = "src/bin/collapse-recursive.rs"
+required-features = ["cli"]
+
 [[bin]]
 name = "inferno-flamegraph"
 path = "src/bin/flamegraph.rs"

src/bin/collapse-recursive.rs

@@ -0,0 +1,40 @@
+use std::io;
+use std::path::PathBuf;
+
+use clap::Parser;
+use inferno::collapse::recursive::{Folder, Options};
+use inferno::collapse::{Collapse, DEFAULT_NTHREADS};
+use once_cell::sync::Lazy;
+
+static NTHREADS: Lazy<String> = Lazy::new(|| DEFAULT_NTHREADS.to_string());
+
+#[derive(Debug, Parser)]
+#[clap(name = "inferno-collapse-recursive", about)]
+struct Opt {
+    /// Number of threads to use
+    #[clap(
+        short = 'n',
+        long = "nthreads",
+        default_value = &**NTHREADS,
+        value_name = "UINT"
+    )]
+    nthreads: usize,
+
+    #[clap(value_name = "PATH")]
+    /// Collapse output file, or STDIN if not specified
+    infile: Option<PathBuf>,
+}
+
+impl Opt {
+    fn into_parts(self) -> (Option<PathBuf>, Options) {
+        let mut options = Options::default();
+        options.nthreads = self.nthreads;
+        (self.infile, options)
+    }
+}
+
+fn main() -> io::Result<()> {
+    let opt = Opt::parse();
+    let (infile, options) = opt.into_parts();
+    Folder::from(options).collapse_file_to_stdout(infile.as_ref())
+}

src/collapse/mod.rs

@@ -35,6 +35,24 @@ pub mod sample;
 ///   [crate-level documentation]: ../../index.html
 pub mod vtune;
 
+/// Collapse direct recursive backtraces.
+///
+/// Post-process a stack list and merge direct recursive calls.
+/// 
+/// For example, collapses
+/// ```text
+/// main;recursive;recursive;recursive;helper 1
+/// ```
+/// into
+/// ```text
+/// main;recursive;helper 1
+/// ```
+///
+/// See the [crate-level documentation] for details.
+///
+///   [crate-level documentation]: ../../index.html
+pub mod recursive;
+
 /// Stack collapsing for the output of the [Visual Studio built-in profiler](https://docs.microsoft.com/en-us/visualstudio/profiling/profiling-feature-tour?view=vs-2019).
 ///
 /// See the [crate-level documentation] for details.

src/collapse/recursive.rs

@@ -0,0 +1,215 @@
+use super::common::{self, CollapsePrivate};
+use std::{borrow::Cow, io};
+
+/// Recursive backtrace folder configuration options.
+#[derive(Clone, Debug)]
+#[non_exhaustive]
+pub struct Options {
+    /// The number of threads to use.
+    ///
+    /// Default is the number of logical cores on your machine.
+    pub nthreads: usize,
+}
+
+impl Default for Options {
+    fn default() -> Self {
+        Self {
+            nthreads: *common::DEFAULT_NTHREADS,
+        }
+    }
+}
+
+/// A "middleware" folder that receives and outputs the folded stack format
+/// expected by [`crate::flamegraph::from_lines`], collapsing direct recursive
+/// backtraces.
+#[derive(Clone)]
+pub struct Folder {
+    /// The number of stacks per job to send to the threadpool.
+    nstacks_per_job: usize,
+
+    // Options...
+    opt: Options,
+}
+
+impl From<Options> for Folder {
+    fn from(mut opt: Options) -> Self {
+        if opt.nthreads == 0 {
+            opt.nthreads = 1;
+        }
+        Self {
+            nstacks_per_job: common::DEFAULT_NSTACKS_PER_JOB,
+            opt,
+        }
+    }
+}
+
+impl Default for Folder {
+    fn default() -> Self {
+        Options::default().into()
+    }
+}
+
+impl CollapsePrivate for Folder {
+    fn pre_process<R>(
+        &mut self,
+        _reader: &mut R,
+        _occurrences: &mut super::common::Occurrences,
+    ) -> std::io::Result<()>
+    where
+        R: std::io::BufRead,
+    {
+        // Don't expect any header.
+        Ok(())
+    }
+
+    fn collapse_single_threaded<R>(
+        &mut self,
+        reader: R,
+        occurrences: &mut super::common::Occurrences,
+    ) -> std::io::Result<()>
+    where
+        R: std::io::BufRead,
+    {
+        for line in reader.lines() {
+            let line = line?;
+            let (stack, count) = Self::line_parts(&line)
+                .ok_or_else(|| io::Error::from(io::ErrorKind::InvalidData))?;
+
+            occurrences.insert_or_add(Self::collapse_stack(stack.into()).into_owned(), count);
+        }
+        Ok(())
+    }
+
+    fn would_end_stack(&mut self, _line: &[u8]) -> bool {
+        // For our purposes, every line is an independent stack
+        true
+    }
+
+    fn clone_and_reset_stack_context(&self) -> Self {
+        self.clone()
+    }
+
+    fn is_applicable(&mut self, _input: &str) -> Option<bool> {
+        // It seems doubtful that the user would ever want to guess to collapse
+        // recursive traces, so let's just never consider ourselves applicable.
+        Some(false)
+    }
+
+    fn nstacks_per_job(&self) -> usize {
+        self.nstacks_per_job
+    }
+
+    fn set_nstacks_per_job(&mut self, n: usize) {
+        self.nstacks_per_job = n;
+    }
+
+    fn nthreads(&self) -> usize {
+        self.opt.nthreads
+    }
+
+    fn set_nthreads(&mut self, n: usize) {
+        self.opt.nthreads = n;
+    }
+}
+
+impl Folder {
+    fn line_parts(line: &str) -> Option<(&str, usize)> {
+        let mut parts = line.rsplitn(2, ' ');
+        let count = parts.next()?.parse().ok()?;
+        let stack = parts.next()?;
+        Some((stack, count))
+    }
+
+    fn collapse_stack(stack: Cow<str>) -> Cow<str> {
+        // First, determine whether we can avoid allocation by just returning
+        // the original stack (in the case that there is no recursion, which is
+        // likely the mainline case).
+        if !Self::is_recursive(&stack) {
+            return stack;
+        }
+
+        // There is recursion, so we can't get away without allocating a new
+        // String.
+        let mut result = String::with_capacity(stack.len());
+        let mut last = None;
+        for frame in stack.split(';') {
+            match last {
+                None => {
+                    result.push_str(frame);
+                    result.push(';')
+                }
+                Some(l) => {
+                    if l != frame {
+                        result.push_str(frame);
+                        result.push(';')
+                    }
+                }
+            }
+            last = Some(frame);
+        }
+
+        // Remove the trailing semicolon
+        result.pop();
+
+        result.into()
+    }
+
+    /// Determine whether or not a stack contains direct recursion.
+    fn is_recursive(stack: &str) -> bool {
+        let mut last = None;
+        for current in stack.split(';') {
+            match last {
+                None => {
+                    last = Some(current);
+                }
+                Some(l) => {
+                    if l == current {
+                        // Recursion!
+                        return true;
+                    } else {
+                        last = Some(current);
+                    }
+                }
+            }
+        }
+        false
+    }
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+
+    #[test]
+    fn test_collapse_stack() {
+        assert_eq!(Folder::collapse_stack("".into()), "");
+        assert_eq!(Folder::collapse_stack("single".into()), "single");
+        assert_eq!(
+            Folder::collapse_stack("not;recursive".into()),
+            "not;recursive"
+        );
+        assert_eq!(
+            Folder::collapse_stack("has;some;some;recursion;recursion".into()),
+            "has;some;recursion"
+        );
+        assert_eq!(
+            Folder::collapse_stack("co;recursive;co;recursive".into()),
+            "co;recursive;co;recursive"
+        );
+    }
+
+    #[test]
+    fn test_line_parts() {
+        assert_eq!(
+            Folder::line_parts("foo;bar;baz 42"),
+            Some(("foo;bar;baz", 42))
+        );
+        assert_eq!(
+            Folder::line_parts("something; including spaces  42"),
+            Some(("something; including spaces ", 42))
+        );
+        assert_eq!(Folder::line_parts(""), None);
+        assert_eq!(Folder::line_parts("no;number"), None);
+    }
+}

tests/collapse-recursive.rs

@@ -0,0 +1,57 @@
+mod common;
+
+use std::fs::File;
+use std::io::{self, BufReader, Cursor};
+use std::process::{Command, Stdio};
+
+use assert_cmd::cargo::CommandCargoExt;
+use inferno::collapse::recursive::{Folder, Options};
+
+fn test_collapse_recursive(
+    test_file: &str,
+    expected_file: &str,
+    options: Options,
+) -> io::Result<()> {
+    for &n in &[1, 2] {
+        let mut options = options.clone();
+        options.nthreads = n;
+        common::test_collapse(Folder::from(options), test_file, expected_file, false)?;
+    }
+    Ok(())
+}
+
+#[test]
+fn collapse_recursive_basic() {
+    let test_file = "./tests/data/collapse-recursive/basic.txt";
+    let result_file = "./tests/data/collapse-recursive/results/basic-collapsed.txt";
+    test_collapse_recursive(test_file, result_file, Options::default()).unwrap()
+}
+
+#[test]
+fn collapse_recursive_cli() {
+    let input_file = "./tests/data/collapse-recursive/basic.txt";
+    let expected_file = "./tests/data/collapse-recursive/results/basic-collapsed.txt";
+
+    // Test with file passed in
+    let output = Command::cargo_bin("inferno-collapse-recursive")
+        .unwrap()
+        .arg(input_file)
+        .output()
+        .expect("failed to execute process");
+    let expected = BufReader::new(File::open(expected_file).unwrap());
+    common::compare_results(Cursor::new(output.stdout), expected, expected_file, false);
+
+    // Test with STDIN
+    let mut child = Command::cargo_bin("inferno-collapse-recursive")
+        .unwrap()
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .spawn()
+        .expect("Failed to spawn child process");
+    let mut input = BufReader::new(File::open(input_file).unwrap());
+    let stdin = child.stdin.as_mut().expect("Failed to open stdin");
+    io::copy(&mut input, stdin).unwrap();
+    let output = child.wait_with_output().expect("Failed to read stdout");
+    let expected = BufReader::new(File::open(expected_file).unwrap());
+    common::compare_results(Cursor::new(output.stdout), expected, expected_file, false);
+}

tests/data/collapse-recursive/basic.txt

@@ -0,0 +1,3 @@
+main;recursive;recursive;recursive;helper 1
+main;recursive;recursive;helper 2
+main;not;recursive 4

tests/data/collapse-recursive/results/basic-collapsed.txt

@@ -0,0 +1,2 @@
+main;not;recursive 4
+main;recursive;helper 3