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
#[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::os::raw::c_int;
use std::path::Path;
pub enum Output<'a> {
#[cfg(feature = "flamegraph")]
Flamegraph(Option<FlamegraphOptions<'a>>),
#[cfg(feature = "protobuf")]
Protobuf,
}
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,
}
}
}
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",
};
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();
profile
.encode(&mut content)
.expect("Error while encoding protobuf");
output_file
.write_all(&content)
.expect("Error while writing protobuf");
}
}
}
}
}