Line data Source code
1 : use anyhow::bail;
2 : use flate2::write::{GzDecoder, GzEncoder};
3 : use flate2::Compression;
4 : use itertools::Itertools as _;
5 : use pprof::protos::{Function, Line, Location, 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(symbol_name) = symbol.name() else {
61 0 : return;
62 : };
63 :
64 0 : let function_name = format!("{symbol_name:#}");
65 0 : let functions_len = functions.len();
66 0 : let function_id = functions
67 0 : .entry(function_name)
68 0 : .or_insert_with_key(|function_name| {
69 0 : let function_id = functions_len as u64 + 1;
70 0 : let system_name = String::from_utf8_lossy(symbol_name.as_bytes());
71 0 : let filename = symbol
72 0 : .filename()
73 0 : .map(|path| path.to_string_lossy())
74 0 : .unwrap_or(Cow::Borrowed(""));
75 0 : Function {
76 0 : id: function_id,
77 0 : name: string_id(function_name),
78 0 : system_name: string_id(&system_name),
79 0 : filename: string_id(&filename),
80 0 : ..Default::default()
81 0 : }
82 0 : })
83 0 : .id;
84 0 : loc.line.push(Line {
85 0 : function_id,
86 0 : line: symbol.lineno().unwrap_or(0) as i64,
87 0 : ..Default::default()
88 0 : });
89 0 : });
90 0 : }
91 :
92 : // Store the resolved functions, and mark the mapping as resolved.
93 0 : profile.function = functions.into_values().sorted_by_key(|f| f.id).collect();
94 0 : profile.string_table = strings
95 0 : .into_iter()
96 0 : .sorted_by_key(|(_, i)| *i)
97 0 : .map(|(s, _)| s)
98 0 : .collect();
99 :
100 0 : for mapping in &mut profile.mapping {
101 0 : mapping.has_functions = true;
102 0 : mapping.has_filenames = true;
103 0 : }
104 :
105 0 : Ok(profile)
106 0 : }
107 :
108 : /// Strips locations (stack frames) matching the given mappings (substring) or function names
109 : /// (regex). The function bool specifies whether child frames should be stripped as well.
110 : ///
111 : /// The string definitions are left behind in the profile for simplicity, to avoid rewriting all
112 : /// string references.
113 0 : pub fn strip_locations(
114 0 : mut profile: Profile,
115 0 : mappings: &[&str],
116 0 : functions: &[(Regex, bool)],
117 0 : ) -> Profile {
118 0 : // Strip mappings.
119 0 : let mut strip_mappings: HashSet<u64> = HashSet::new();
120 0 :
121 0 : profile.mapping.retain(|mapping| {
122 0 : let Some(name) = profile.string_table.get(mapping.filename as usize) else {
123 0 : return true;
124 : };
125 0 : if mappings.iter().any(|substr| name.contains(substr)) {
126 0 : strip_mappings.insert(mapping.id);
127 0 : return false;
128 0 : }
129 0 : true
130 0 : });
131 0 :
132 0 : // Strip functions.
133 0 : let mut strip_functions: HashMap<u64, bool> = HashMap::new();
134 0 :
135 0 : profile.function.retain(|function| {
136 0 : let Some(name) = profile.string_table.get(function.name as usize) else {
137 0 : return true;
138 : };
139 0 : for (regex, strip_children) in functions {
140 0 : if regex.is_match(name) {
141 0 : strip_functions.insert(function.id, *strip_children);
142 0 : return false;
143 0 : }
144 : }
145 0 : true
146 0 : });
147 0 :
148 0 : // Strip locations. The bool specifies whether child frames should be stripped too.
149 0 : let mut strip_locations: HashMap<u64, bool> = HashMap::new();
150 0 :
151 0 : profile.location.retain(|location| {
152 0 : for line in &location.line {
153 0 : if let Some(strip_children) = strip_functions.get(&line.function_id) {
154 0 : strip_locations.insert(location.id, *strip_children);
155 0 : return false;
156 0 : }
157 : }
158 0 : if strip_mappings.contains(&location.mapping_id) {
159 0 : strip_locations.insert(location.id, false);
160 0 : return false;
161 0 : }
162 0 : true
163 0 : });
164 :
165 : // Strip sample locations.
166 0 : for sample in &mut profile.sample {
167 : // First, find the uppermost function with child removal and truncate the stack.
168 0 : if let Some(truncate) = sample
169 0 : .location_id
170 0 : .iter()
171 0 : .rposition(|id| strip_locations.get(id) == Some(&true))
172 0 : {
173 0 : sample.location_id.drain(..=truncate);
174 0 : }
175 : // Next, strip any individual frames without child removal.
176 0 : sample
177 0 : .location_id
178 0 : .retain(|id| !strip_locations.contains_key(id));
179 0 : }
180 :
181 0 : profile
182 0 : }
183 :
184 : /// Generates an SVG flamegraph from a symbolized pprof profile.
185 0 : pub fn flamegraph(
186 0 : profile: Profile,
187 0 : opts: &mut inferno::flamegraph::Options,
188 0 : ) -> anyhow::Result<Vec<u8>> {
189 0 : if profile.mapping.iter().any(|m| !m.has_functions) {
190 0 : bail!("profile not symbolized");
191 0 : }
192 0 :
193 0 : // Index locations, functions, and strings.
194 0 : let locations: HashMap<u64, Location> =
195 0 : profile.location.into_iter().map(|l| (l.id, l)).collect();
196 0 : let functions: HashMap<u64, Function> =
197 0 : profile.function.into_iter().map(|f| (f.id, f)).collect();
198 0 : let strings = profile.string_table;
199 0 :
200 0 : // Resolve stacks as function names, and sum sample values per stack. Also reverse the stack,
201 0 : // since inferno expects it bottom-up.
202 0 : let mut stacks: HashMap<Vec<&str>, i64> = HashMap::new();
203 0 : for sample in profile.sample {
204 0 : let mut stack = Vec::with_capacity(sample.location_id.len());
205 0 : for location in sample.location_id.into_iter().rev() {
206 0 : let Some(location) = locations.get(&location) else {
207 0 : bail!("missing location {location}");
208 : };
209 0 : for line in location.line.iter().rev() {
210 0 : let Some(function) = functions.get(&line.function_id) else {
211 0 : bail!("missing function {}", line.function_id);
212 : };
213 0 : let Some(name) = strings.get(function.name as usize) else {
214 0 : bail!("missing string {}", function.name);
215 : };
216 0 : stack.push(name.as_str());
217 : }
218 : }
219 0 : let Some(&value) = sample.value.first() else {
220 0 : bail!("missing value");
221 : };
222 0 : *stacks.entry(stack).or_default() += value;
223 : }
224 :
225 : // Construct stack lines for inferno.
226 0 : let lines = stacks
227 0 : .into_iter()
228 0 : .map(|(stack, value)| (stack.into_iter().join(";"), value))
229 0 : .map(|(stack, value)| format!("{stack} {value}"))
230 0 : .sorted()
231 0 : .collect_vec();
232 0 :
233 0 : // Construct the flamegraph.
234 0 : let mut bytes = Vec::new();
235 0 : let lines = lines.iter().map(|line| line.as_str());
236 0 : inferno::flamegraph::from_lines(opts, lines, &mut bytes)?;
237 0 : Ok(bytes)
238 0 : }
|
