pyroscope_rbspy_oncpu/ui/
flamegraph.rs

1use anyhow::Result;
2use inferno::flamegraph::{Direction, Options};
3use std::collections::HashMap;
4use std::io::Write;
5
6use crate::core::types::StackFrame;
7
8// Simple counter that maps stacks to flamegraph collapsed format
9#[derive(Default)]
10pub struct Stats {
11    pub counts: HashMap<String, usize>,
12}
13
14impl Stats {
15    pub fn record(&mut self, stack: &[StackFrame]) -> Result<()> {
16        let frame = stack
17            .iter()
18            .rev()
19            .map(|frame| format!("{}", frame))
20            .collect::<Vec<String>>()
21            .join(";");
22
23        *self.counts.entry(frame).or_insert(0) += 1;
24        Ok(())
25    }
26
27    pub fn write_flamegraph<W: Write>(&self, w: W, min_width: f64) -> Result<()> {
28        if self.is_empty() {
29            eprintln!("Warning: no profile samples were collected");
30        } else {
31            let mut opts = Options::default();
32            opts.direction = Direction::Inverted;
33            opts.hash = true;
34            opts.min_width = min_width;
35            inferno::flamegraph::from_lines(
36                &mut opts,
37                self.get_lines().iter().map(|x| x.as_str()),
38                w,
39            )?;
40        }
41
42        Ok(())
43    }
44
45    pub fn write_collapsed<W: Write>(&self, w: &mut W) -> Result<()> {
46        if self.is_empty() {
47            eprintln!("Warning: no profile samples were collected");
48        } else {
49            self.get_lines()
50                .iter()
51                .try_for_each(|line| write!(w, "{}\n", line))?;
52        }
53        Ok(())
54    }
55
56    fn get_lines(&self) -> Vec<String> {
57        self.counts
58            .iter()
59            .map(|(frame, count)| format!("{} {}", frame, count))
60            .collect()
61    }
62
63    fn is_empty(&self) -> bool {
64        self.counts.is_empty()
65    }
66}
67
68#[cfg(test)]
69mod tests {
70    use crate::ui::flamegraph::*;
71    use std::io::Cursor;
72
73    // Build a test stackframe
74    fn f(i: usize) -> StackFrame {
75        StackFrame {
76            name: format!("func{}", i),
77            relative_path: format!("file{}.rb", i),
78            absolute_path: None,
79            lineno: Some(i),
80        }
81    }
82
83    // Build test stats
84    fn build_stats() -> Result<Stats> {
85        let mut stats = Stats::default();
86        stats.record(&vec![f(1)])?;
87        stats.record(&vec![f(2), f(1)])?;
88        stats.record(&vec![f(2), f(1)])?;
89        stats.record(&vec![f(2), f(3), f(1)])?;
90        stats.record(&vec![f(2), f(3), f(1)])?;
91        stats.record(&vec![f(2), f(3), f(1)])?;
92        Ok(stats)
93    }
94
95    fn assert_contains(counts: &HashMap<String, usize>, s: &str, val: usize) {
96        assert_eq!(counts.get(&s.to_string()), Some(&val));
97    }
98
99    #[test]
100    fn test_stats() -> Result<()> {
101        let stats = build_stats()?;
102        let counts = &stats.counts;
103        assert_contains(counts, "func1 - file1.rb:1", 1);
104        assert_contains(
105            counts,
106            "func1 - file1.rb:1;func3 - file3.rb:3;func2 - file2.rb:2",
107            3,
108        );
109        assert_contains(counts, "func1 - file1.rb:1;func2 - file2.rb:2", 2);
110
111        Ok(())
112    }
113
114    #[test]
115    fn test_collapsed() -> Result<()> {
116        let stats = build_stats()?;
117        let mut writer = Cursor::new(Vec::<u8>::new());
118        stats.write_collapsed(&mut writer)?;
119        let collapsed_text = std::str::from_utf8(writer.get_ref())?;
120        assert!(collapsed_text.contains("func1 - file1.rb:1 1"));
121        assert!(
122            collapsed_text.contains("func1 - file1.rb:1;func3 - file3.rb:3;func2 - file2.rb:2 3")
123        );
124        assert!(collapsed_text.contains("func1 - file1.rb:1;func2 - file2.rb:2 2"));
125
126        Ok(())
127    }
128
129    #[test]
130    fn test_flamegraph_from_collapsed() -> Result<()> {
131        let stats = build_stats()?;
132
133        let mut writer = Cursor::new(Vec::<u8>::new());
134        stats.write_collapsed(&mut writer)?;
135
136        let collapsed_reader = Cursor::new(writer.into_inner());
137        let svg_writer = Cursor::new(Vec::new());
138
139        let mut opts = Options::default();
140        opts.direction = Direction::Inverted;
141        opts.min_width = 0.1;
142        inferno::flamegraph::from_reader(&mut opts, collapsed_reader, svg_writer)?;
143
144        Ok(())
145    }
146}
pyroscope_rbspy_oncpu/ui/flamegraph.rs

pyroscope_rbspy_oncpu/ui/
flamegraph.rs
