lance_testing/
pprof.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{fs::File, os::raw::c_int, path::Path};
5
6use ::pprof::{ProfilerGuard, flamegraph::Options as FlamegraphOptions};
7use criterion::profiler::Profiler;
8
9#[allow(clippy::large_enum_variant)]
10pub enum Output<'a> {
11    Flamegraph(Option<FlamegraphOptions<'a>>),
12}
13
14pub struct PProfProfiler<'a, 'b> {
15    frequency: c_int,
16    output: Output<'b>,
17    active_profiler: Option<ProfilerGuard<'a>>,
18}
19
20impl<'a, 'b> PProfProfiler<'a, 'b> {
21    pub fn new(frequency: c_int, output: Output<'b>) -> Self {
22        Self {
23            frequency,
24            output,
25            active_profiler: None,
26        }
27    }
28}
29
30impl<'a, 'b> Profiler for PProfProfiler<'a, 'b> {
31    fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) {
32        self.active_profiler = Some(ProfilerGuard::new(self.frequency).unwrap());
33    }
34
35    fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
36        std::fs::create_dir_all(benchmark_dir).unwrap();
37
38        let Output::Flamegraph(options) = &mut self.output;
39
40        let output_path = benchmark_dir.join("flamegraph.svg");
41        let output_file = File::create(&output_path).unwrap_or_else(|_| {
42            panic!("File system error while creating {}", output_path.display())
43        });
44
45        if let Some(profiler) = self.active_profiler.take() {
46            let default_options = &mut FlamegraphOptions::default();
47            let options = options.as_mut().unwrap_or(default_options);
48
49            profiler
50                .report()
51                .build()
52                .unwrap()
53                .flamegraph_with_options(output_file, options)
54                .expect("Error while writing flamegraph");
55        }
56    }
57}
lance_testing/pprof.rs

lance_testing/
pprof.rs
