Skip to content

Commit 413cd90

Browse files
committed
feat: implement Report::from_pprof() for protobuf roundtrip conversion
Add backwards conversion from protobuf Profile to Report, enabling roundtrip conversion between the two formats. This allows users to reconstruct Report objects from previously exported protobuf data, which is helpful for making flamegraphs. Changes: - Add Report::from_pprof() method to convert protobuf Profile back to Report - Fix pprof() method to create one location per frame instead of per symbol - Group inlined functions correctly within locations per pprof spec - Add roundtrip conversion test The changes make the protobuf output more standards-compliant by grouping inlined function symbols within the same location, which aligns with how Google's pprof tool expects the data to be structured.
1 parent 0c52c5f commit 413cd90

File tree

1 file changed

+211
-25
lines changed

1 file changed

+211
-25
lines changed

src/report.rs

Lines changed: 211 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,110 @@ mod protobuf {
226226
const THREAD: &str = "thread";
227227

228228
impl Report {
229+
/// Create a `Report` from a protobuf `Profile`. This can be useful
230+
/// for creating a flamegraph from a saved protobuf file.
231+
pub fn from_pprof(profile: &protos::Profile) -> crate::Result<Self> {
232+
let mut data = HashMap::new();
233+
234+
let strings: Vec<&str> = profile.string_table.iter().map(|s| s.as_str()).collect();
235+
236+
let mut functions = HashMap::new();
237+
for func in profile.function.iter() {
238+
functions.insert(func.id, func);
239+
}
240+
241+
let mut locations = HashMap::new();
242+
for loc in profile.location.iter() {
243+
locations.insert(loc.id, loc);
244+
}
245+
246+
for sample in profile.sample.iter() {
247+
let mut frames = Vec::new();
248+
249+
for &loc_id in sample.location_id.iter() {
250+
if let Some(location) = locations.get(&loc_id) {
251+
let mut symbols = Vec::new();
252+
253+
for line in location.line.iter() {
254+
if let Some(function) = functions.get(&line.function_id) {
255+
let name =
256+
strings.get(function.name as usize).unwrap_or(&"Unknown");
257+
let filename = strings
258+
.get(function.filename as usize)
259+
.unwrap_or(&"Unknown");
260+
261+
let symbol = crate::Symbol {
262+
name: Some(name.as_bytes().to_vec()),
263+
addr: None,
264+
lineno: if line.line > 0 {
265+
Some(line.line as u32)
266+
} else {
267+
None
268+
},
269+
filename: if *filename != "Unknown" {
270+
Some(filename.into())
271+
} else {
272+
None
273+
},
274+
};
275+
symbols.push(symbol);
276+
}
277+
}
278+
279+
if !symbols.is_empty() {
280+
frames.push(symbols);
281+
}
282+
}
283+
}
284+
285+
// Extract thread name from labels
286+
let mut thread_name = String::new();
287+
for label in sample.label.iter() {
288+
let key_str = strings.get(label.key as usize).unwrap_or(&"");
289+
if *key_str == THREAD {
290+
thread_name = strings.get(label.str as usize).unwrap_or(&"").to_string();
291+
break;
292+
}
293+
}
294+
295+
let frames_key = Frames {
296+
frames,
297+
thread_name,
298+
thread_id: 0, // Not preserved in protobuf format
299+
sample_timestamp: SystemTime::UNIX_EPOCH, // Not preserved
300+
};
301+
302+
let count = sample.value.first().copied().unwrap_or(0) as isize;
303+
*data.entry(frames_key).or_insert(0) += count;
304+
}
305+
306+
let frequency = if profile.period > 0 {
307+
(1_000_000_000 / profile.period) as i32
308+
} else {
309+
1
310+
};
311+
312+
let start_time = if profile.time_nanos > 0 {
313+
SystemTime::UNIX_EPOCH + std::time::Duration::from_nanos(profile.time_nanos as u64)
314+
} else {
315+
SystemTime::UNIX_EPOCH
316+
};
317+
318+
let duration = if profile.duration_nanos > 0 {
319+
std::time::Duration::from_nanos(profile.duration_nanos as u64)
320+
} else {
321+
std::time::Duration::default()
322+
};
323+
324+
let timing = crate::timer::ReportTiming {
325+
frequency,
326+
start_time,
327+
duration,
328+
};
329+
330+
Ok(Report { data, timing })
331+
}
332+
229333
/// `pprof` will generate google's pprof format report.
230334
pub fn pprof(&self) -> crate::Result<protos::Profile> {
231335
let mut dedup_str = HashSet::new();
@@ -260,40 +364,45 @@ mod protobuf {
260364
for (key, count) in self.data.iter() {
261365
let mut locs = vec![];
262366
for frame in key.frames.iter() {
367+
let location_id = loc_tbl.len() as u64 + 1;
368+
let mut lines = vec![];
369+
263370
for symbol in frame {
264371
let name = symbol.name();
265-
if let Some(loc_idx) = functions.get(&name) {
266-
locs.push(*loc_idx);
267-
continue;
268-
}
269-
let sys_name = symbol.sys_name();
270-
let filename = symbol.filename();
271-
let lineno = symbol.lineno();
272-
let function_id = fn_tbl.len() as u64 + 1;
273-
let function = protos::Function {
274-
id: function_id,
275-
name: *strings.get(name.as_str()).unwrap() as i64,
276-
system_name: *strings.get(sys_name.as_ref()).unwrap() as i64,
277-
filename: *strings.get(filename.as_ref()).unwrap() as i64,
278-
..protos::Function::default()
372+
let function_id = if let Some(&existing_id) = functions.get(&name) {
373+
existing_id
374+
} else {
375+
let sys_name = symbol.sys_name();
376+
let filename = symbol.filename();
377+
let function_id = fn_tbl.len() as u64 + 1;
378+
let function = protos::Function {
379+
id: function_id,
380+
name: *strings.get(name.as_str()).unwrap() as i64,
381+
system_name: *strings.get(sys_name.as_ref()).unwrap() as i64,
382+
filename: *strings.get(filename.as_ref()).unwrap() as i64,
383+
..protos::Function::default()
384+
};
385+
functions.insert(name, function_id);
386+
fn_tbl.push(function);
387+
function_id
279388
};
280-
functions.insert(name, function_id);
389+
390+
let lineno = symbol.lineno();
281391
let line = protos::Line {
282392
function_id,
283393
line: lineno as i64,
284394
..protos::Line::default()
285395
};
286-
let loc = protos::Location {
287-
id: function_id,
288-
line: vec![line].into(),
289-
..protos::Location::default()
290-
};
291-
// the fn_tbl has the same length with loc_tbl
292-
fn_tbl.push(function);
293-
loc_tbl.push(loc);
294-
// current frame locations
295-
locs.push(function_id);
396+
lines.push(line);
296397
}
398+
399+
let loc = protos::Location {
400+
id: location_id,
401+
line: lines.into(),
402+
..protos::Location::default()
403+
};
404+
loc_tbl.push(loc);
405+
locs.push(location_id);
297406
}
298407
let thread_name = protos::Label {
299408
key: *strings.get(THREAD).unwrap() as i64,
@@ -341,4 +450,81 @@ mod protobuf {
341450
Ok(profile)
342451
}
343452
}
453+
454+
#[cfg(test)]
455+
mod tests {
456+
use super::*;
457+
use std::collections::HashSet;
458+
459+
#[test]
460+
fn test_roundtrip_conversion() {
461+
let guard = crate::ProfilerGuard::new(100).unwrap();
462+
463+
// Generate profiling data with different call patterns
464+
for i in 0..100000 {
465+
if i % 3 == 0 {
466+
expensive_function_a(i);
467+
} else if i % 3 == 1 {
468+
expensive_function_b(i);
469+
} else {
470+
expensive_function_c(i);
471+
}
472+
}
473+
474+
let report = guard.report().build().unwrap();
475+
assert!(
476+
!report.data.is_empty(),
477+
"Should have captured some profiling data"
478+
);
479+
480+
let profile = report.pprof().unwrap();
481+
let restored_report = Report::from_pprof(&profile).unwrap();
482+
483+
let original_symbols: HashSet<String> = report
484+
.data
485+
.keys()
486+
.flat_map(|frames| frames.frames.iter())
487+
.flat_map(|frame| frame.iter())
488+
.map(|symbol| symbol.name())
489+
.collect();
490+
491+
let restored_symbols: HashSet<String> = restored_report
492+
.data
493+
.keys()
494+
.flat_map(|frames| frames.frames.iter())
495+
.flat_map(|frame| frame.iter())
496+
.map(|symbol| symbol.name())
497+
.collect();
498+
499+
assert_eq!(original_symbols.len(), restored_symbols.len());
500+
for symbol in &original_symbols {
501+
assert!(restored_symbols.contains(symbol));
502+
}
503+
504+
let original_total: isize = report.data.values().sum();
505+
let restored_total: isize = restored_report.data.values().sum();
506+
assert_eq!(original_total, restored_total);
507+
508+
assert_eq!(report.timing.frequency, restored_report.timing.frequency);
509+
}
510+
511+
#[inline(never)]
512+
fn expensive_function_a(n: usize) -> usize {
513+
(0..n % 100).map(|i| i * i).sum()
514+
}
515+
516+
#[inline(never)]
517+
fn expensive_function_b(n: usize) -> usize {
518+
(0..n % 50).fold(1, |acc, x| acc.wrapping_mul(x + 1))
519+
}
520+
521+
#[inline(never)]
522+
fn expensive_function_c(n: usize) -> usize {
523+
let mut result = n;
524+
for i in 0..n % 30 {
525+
result = result.wrapping_add(i * 3);
526+
}
527+
result
528+
}
529+
}
344530
}

0 commit comments

Comments
 (0)