trace_tools/util/
flamegrapher.rs

1// Copyright 2022 IOTA Stiftung
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::error::{Error, FlamegrapherErrorKind};
5
6use std::{
7    fs::File,
8    io::{BufReader, BufWriter},
9    path::{Path, PathBuf},
10};
11
12/// Helper struct that uses [`inferno`] internally to programatically produce a flamegraph from a folded
13/// stack trace file.
14#[derive(Default)]
15pub struct Flamegrapher {
16    stack_filename: Option<PathBuf>,
17    graph_filename: Option<PathBuf>,
18}
19
20impl Flamegrapher {
21    /// Creates a new [`Flamegrapher`] with no associated files.
22    pub fn new() -> Self {
23        Self::default()
24    }
25
26    /// Returns a [`Flamegrapher`] with the given folded stack file name.
27    ///
28    /// # Errors
29    /// This method may fail in the following ways:
30    ///  - The given stack file does not exist.
31    pub fn with_stack_file<P: AsRef<Path>>(mut self, stack_filename: P) -> Result<Self, Error> {
32        let stack_filename = stack_filename.as_ref().to_path_buf();
33
34        if !stack_filename.exists() {
35            return Err(Error::Flamegrapher(FlamegrapherErrorKind::StackFileNotFound(
36                stack_filename,
37            )));
38        }
39
40        self.stack_filename = Some(stack_filename);
41        Ok(self)
42    }
43
44    /// Returns a [`Flamegrapher`] with the given flamegraph file name.
45    ///
46    /// This file will be given the extension of `.svg`.
47    ///
48    /// # Errors
49    /// This method may fail in the following ways:
50    ///  - The given graph filename is invalid (it does not belong to a directory that exists).
51    pub fn with_graph_file<P: AsRef<Path>>(mut self, graph_filename: P) -> Result<Self, Error> {
52        let graph_filename = graph_filename.as_ref().with_extension("svg");
53
54        match graph_filename.parent() {
55            Some(directory) if !directory.is_dir() => {
56                return Err(Error::Flamegrapher(FlamegrapherErrorKind::GraphFileInvalid(
57                    graph_filename,
58                )));
59            }
60            _ => {}
61        }
62
63        self.graph_filename = Some(graph_filename);
64        Ok(self)
65    }
66
67    /// Uses [`inferno`] to generate a flamegraph from the given folded stack file, and writes it to the given
68    /// output image file.
69    ///
70    /// # Errors
71    /// This method may fail in the following ways:
72    ///  - This [`Flamegrapher`] does not have a stack or graph file associated with it.
73    ///  - An error was encountered when opening the folded stack file for reading.
74    ///  - An error was encountered when creating the graph file.
75    pub fn write_flamegraph(&self) -> Result<(), Error> {
76        let stack_filename = self
77            .stack_filename
78            .as_ref()
79            .ok_or_else(|| Error::Flamegrapher(FlamegrapherErrorKind::MissingField("stack_filename".to_string())))?;
80
81        let graph_filename = self
82            .graph_filename
83            .as_ref()
84            .ok_or_else(|| Error::Flamegrapher(FlamegrapherErrorKind::MissingField("graph_filename".to_string())))?;
85
86        let stack_file = File::open(stack_filename).map_err(|err| Error::Flamegrapher(err.into()))?;
87        let reader = BufReader::new(stack_file);
88
89        let graph_file = File::create(graph_filename).map_err(|err| Error::Flamegrapher(err.into()))?;
90        let writer = BufWriter::new(graph_file);
91
92        let mut graph_options = inferno::flamegraph::Options::default();
93        inferno::flamegraph::from_reader(&mut graph_options, reader, writer)
94            .map_err(|err| Error::Flamegrapher(FlamegrapherErrorKind::Inferno(Box::new(err))))?;
95
96        Ok(())
97    }
98}
trace_tools/util/flamegrapher.rs

trace_tools/util/
flamegrapher.rs
