1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

//! A tracing flamegraph component.

use color_eyre::eyre::Report;
use std::{
    fs::File,
    io::{BufReader, BufWriter},
    path::Path,
    path::PathBuf,
};
use tracing::Subscriber;
use tracing_subscriber::{registry::LookupSpan, Layer};

pub struct Grapher {
    guard: tracing_flame::FlushGuard<BufWriter<File>>,
    path: PathBuf,
}

pub fn layer<S>(path_root: &Path) -> (impl Layer<S>, Grapher)
where
    S: Subscriber + for<'span> LookupSpan<'span>,
{
    let path = path_root.with_extension("folded");
    let (layer, guard) = tracing_flame::FlameLayer::with_file(&path).unwrap();
    let layer = layer.with_empty_samples(false).with_threads_collapsed(true);
    let flamegrapher = Grapher { guard, path };
    (layer, flamegrapher)
}

impl Grapher {
    pub fn write_flamegraph(&self) -> Result<(), Report> {
        self.guard.flush()?;
        let out_path = self.path.with_extension("svg");
        let inf = File::open(&self.path)?;
        let reader = BufReader::new(inf);

        let out = File::create(out_path)?;
        let writer = BufWriter::new(out);

        let mut opts = inferno::flamegraph::Options::default();
        debug!("writing flamegraph to disk...");
        inferno::flamegraph::from_reader(&mut opts, reader, writer)?;

        Ok(())
    }
}