@@ -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