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
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

// Copyright 2020 TiKV Project Authors. Licensed under Apache-2.0.

#[cfg(feature = "flamegraph")]
use crate::flamegraph::Options as FlamegraphOptions;
#[cfg(feature = "_protobuf")]
use crate::protos::Message;

use crate::ProfilerGuard;
use criterion::profiler::Profiler;

use std::fs::File;
#[cfg(feature = "_protobuf")]
use std::io::Write;
use std::marker::PhantomData;
use std::os::raw::c_int;
use std::path::Path;

#[allow(clippy::large_enum_variant)]
pub enum Output<'a> {
    #[cfg(feature = "flamegraph")]
    Flamegraph(Option<FlamegraphOptions<'a>>),

    #[cfg(feature = "_protobuf")]
    Protobuf,

    #[deprecated(
        note = "This branch is used to include lifetime parameter. Don't use it directly."
    )]
    _Phantom(PhantomData<&'a ()>),
}

pub struct PProfProfiler<'a, 'b> {
    frequency: c_int,
    output: Output<'b>,
    active_profiler: Option<ProfilerGuard<'a>>,
}

impl<'a, 'b> PProfProfiler<'a, 'b> {
    pub fn new(frequency: c_int, output: Output<'b>) -> Self {
        Self {
            frequency,
            output,
            active_profiler: None,
        }
    }
}

#[cfg(not(any(feature = "_protobuf", feature = "flamegraph")))]
compile_error!("Either feature \"protobuf\" or \"flamegraph\" must be enabled when \"criterion\" feature is enabled.");

impl<'a, 'b> Profiler for PProfProfiler<'a, 'b> {
    fn start_profiling(&mut self, _benchmark_id: &str, _benchmark_dir: &Path) {
        self.active_profiler = Some(ProfilerGuard::new(self.frequency).unwrap());
    }

    fn stop_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
        std::fs::create_dir_all(benchmark_dir).unwrap();

        let filename = match self.output {
            #[cfg(feature = "flamegraph")]
            Output::Flamegraph(_) => "flamegraph.svg",
            #[cfg(feature = "_protobuf")]
            Output::Protobuf => "profile.pb",
            // This is `""` but not `unreachable!()`, because `unreachable!()`
            // will result in another compile error, so that the user may not
            // realize the error thrown by `compile_error!()` at the first time.
            _ => "",
        };
        let output_path = benchmark_dir.join(filename);
        let output_file = File::create(&output_path).unwrap_or_else(|_| {
            panic!("File system error while creating {}", output_path.display())
        });

        if let Some(profiler) = self.active_profiler.take() {
            match &mut self.output {
                #[cfg(feature = "flamegraph")]
                Output::Flamegraph(options) => {
                    let default_options = &mut FlamegraphOptions::default();
                    let options = options.as_mut().unwrap_or(default_options);

                    profiler
                        .report()
                        .build()
                        .unwrap()
                        .flamegraph_with_options(output_file, options)
                        .expect("Error while writing flamegraph");
                }

                #[cfg(feature = "_protobuf")]
                Output::Protobuf => {
                    let mut output_file = output_file;

                    let profile = profiler.report().build().unwrap().pprof().unwrap();

                    let mut content = Vec::new();
                    #[cfg(not(feature = "protobuf-codec"))]
                    profile
                        .encode(&mut content)
                        .expect("Error while encoding protobuf");
                    #[cfg(feature = "protobuf-codec")]
                    profile
                        .write_to_vec(&mut content)
                        .expect("Error while encoding protobuf");

                    output_file
                        .write_all(&content)
                        .expect("Error while writing protobuf");
                }

                _ => unreachable!(),
            }
        }
    }
}