Line data Source code
1 : use flate2::write::{GzDecoder, GzEncoder};
2 : use flate2::Compression;
3 : use itertools::Itertools as _;
4 : use once_cell::sync::Lazy;
5 : use pprof::protos::{Function, Line, Message as _, Profile};
6 : use regex::Regex;
7 :
8 : use std::borrow::Cow;
9 : use std::collections::{HashMap, HashSet};
10 : use std::ffi::c_void;
11 : use std::io::Write as _;
12 :
13 : /// Decodes a gzip-compressed Protobuf-encoded pprof profile.
14 0 : pub fn decode(bytes: &[u8]) -> anyhow::Result<Profile> {
15 0 : let mut gz = GzDecoder::new(Vec::new());
16 0 : gz.write_all(bytes)?;
17 0 : Ok(Profile::parse_from_bytes(&gz.finish()?)?)
18 0 : }
19 :
20 : /// Encodes a pprof profile as gzip-compressed Protobuf.
21 0 : pub fn encode(profile: &Profile) -> anyhow::Result<Vec<u8>> {
22 0 : let mut gz = GzEncoder::new(Vec::new(), Compression::default());
23 0 : profile.write_to_writer(&mut gz)?;
24 0 : Ok(gz.finish()?)
25 0 : }
26 :
27 : /// Symbolizes a pprof profile using the current binary.
28 0 : pub fn symbolize(mut profile: Profile) -> anyhow::Result<Profile> {
29 0 : if !profile.function.is_empty() {
30 0 : return Ok(profile); // already symbolized
31 0 : }
32 0 :
33 0 : // Collect function names.
34 0 : let mut functions: HashMap<String, Function> = HashMap::new();
35 0 : let mut strings: HashMap<String, i64> = profile
36 0 : .string_table
37 0 : .into_iter()
38 0 : .enumerate()
39 0 : .map(|(i, s)| (s, i as i64))
40 0 : .collect();
41 0 :
42 0 : // Helper to look up or register a string.
43 0 : let mut string_id = |s: &str| -> i64 {
44 : // Don't use .entry() to avoid unnecessary allocations.
45 0 : if let Some(id) = strings.get(s) {
46 0 : return *id;
47 0 : }
48 0 : let id = strings.len() as i64;
49 0 : strings.insert(s.to_string(), id);
50 0 : id
51 0 : };
52 :
53 0 : for loc in &mut profile.location {
54 0 : if !loc.line.is_empty() {
55 0 : continue;
56 0 : }
57 0 :
58 0 : // Resolve the line and function for each location.
59 0 : backtrace::resolve(loc.address as *mut c_void, |symbol| {
60 0 : let Some(symname) = symbol.name() else {
61 0 : return;
62 : };
63 0 : let mut name = symname.to_string();
64 :
65 : // Strip the Rust monomorphization suffix from the symbol name.
66 : static SUFFIX_REGEX: Lazy<Regex> =
67 0 : Lazy::new(|| Regex::new("::h[0-9a-f]{16}$").expect("invalid regex"));
68 0 : if let Some(m) = SUFFIX_REGEX.find(&name) {
69 0 : name.truncate(m.start());
70 0 : }
71 :
72 0 : let function_id = match functions.get(&name) {
73 0 : Some(function) => function.id,
74 : None => {
75 0 : let id = functions.len() as u64 + 1;
76 0 : let system_name = String::from_utf8_lossy(symname.as_bytes());
77 0 : let filename = symbol
78 0 : .filename()
79 0 : .map(|path| path.to_string_lossy())
80 0 : .unwrap_or(Cow::Borrowed(""));
81 0 : let function = Function {
82 0 : id,
83 0 : name: string_id(&name),
84 0 : system_name: string_id(&system_name),
85 0 : filename: string_id(&filename),
86 0 : ..Default::default()
87 0 : };
88 0 : functions.insert(name, function);
89 0 : id
90 : }
91 : };
92 0 : loc.line.push(Line {
93 0 : function_id,
94 0 : line: symbol.lineno().unwrap_or(0) as i64,
95 0 : ..Default::default()
96 0 : });
97 0 : });
98 0 : }
99 :
100 : // Store the resolved functions, and mark the mapping as resolved.
101 0 : profile.function = functions.into_values().sorted_by_key(|f| f.id).collect();
102 0 : profile.string_table = strings
103 0 : .into_iter()
104 0 : .sorted_by_key(|(_, i)| *i)
105 0 : .map(|(s, _)| s)
106 0 : .collect();
107 :
108 0 : for mapping in &mut profile.mapping {
109 0 : mapping.has_functions = true;
110 0 : mapping.has_filenames = true;
111 0 : }
112 :
113 0 : Ok(profile)
114 0 : }
115 :
116 : /// Strips locations (stack frames) matching the given mappings (substring) or function names
117 : /// (regex). The function bool specifies whether child frames should be stripped as well.
118 : ///
119 : /// The string definitions are left behind in the profile for simplicity, to avoid rewriting all
120 : /// string references.
121 0 : pub fn strip_locations(
122 0 : mut profile: Profile,
123 0 : mappings: &[&str],
124 0 : functions: &[(Regex, bool)],
125 0 : ) -> Profile {
126 0 : // Strip mappings.
127 0 : let mut strip_mappings: HashSet<u64> = HashSet::new();
128 0 :
129 0 : profile.mapping.retain(|mapping| {
130 0 : let Some(name) = profile.string_table.get(mapping.filename as usize) else {
131 0 : return true;
132 : };
133 0 : if mappings.iter().any(|substr| name.contains(substr)) {
134 0 : strip_mappings.insert(mapping.id);
135 0 : return false;
136 0 : }
137 0 : true
138 0 : });
139 0 :
140 0 : // Strip functions.
141 0 : let mut strip_functions: HashMap<u64, bool> = HashMap::new();
142 0 :
143 0 : profile.function.retain(|function| {
144 0 : let Some(name) = profile.string_table.get(function.name as usize) else {
145 0 : return true;
146 : };
147 0 : for (regex, strip_children) in functions {
148 0 : if regex.is_match(name) {
149 0 : strip_functions.insert(function.id, *strip_children);
150 0 : return false;
151 0 : }
152 : }
153 0 : true
154 0 : });
155 0 :
156 0 : // Strip locations. The bool specifies whether child frames should be stripped too.
157 0 : let mut strip_locations: HashMap<u64, bool> = HashMap::new();
158 0 :
159 0 : profile.location.retain(|location| {
160 0 : for line in &location.line {
161 0 : if let Some(strip_children) = strip_functions.get(&line.function_id) {
162 0 : strip_locations.insert(location.id, *strip_children);
163 0 : return false;
164 0 : }
165 : }
166 0 : if strip_mappings.contains(&location.mapping_id) {
167 0 : strip_locations.insert(location.id, false);
168 0 : return false;
169 0 : }
170 0 : true
171 0 : });
172 :
173 : // Strip sample locations.
174 0 : for sample in &mut profile.sample {
175 : // First, find the uppermost function with child removal and truncate the stack.
176 0 : if let Some(truncate) = sample
177 0 : .location_id
178 0 : .iter()
179 0 : .rposition(|id| strip_locations.get(id) == Some(&true))
180 0 : {
181 0 : sample.location_id.drain(..=truncate);
182 0 : }
183 : // Next, strip any individual frames without child removal.
184 0 : sample
185 0 : .location_id
186 0 : .retain(|id| !strip_locations.contains_key(id));
187 0 : }
188 :
189 0 : profile
190 0 : }
|
