Add aggregated graph for instruction tracing

Author: jasonz-dfinity

canbench-bin/src/instruction_tracing.rs

@@ -307,65 +307,16 @@ fn inject_tracing(
 /// Renders the tracing to a file. Adapted from
 /// https://github.com/dfinity/ic-repl/blob/master/src/tracing.rs
 pub(super) fn write_traces_to_file(
-    input: Vec<(i32, i64)>,
+    input: InstructionTraceGraphNode,
     names: &BTreeMap<i32, String>,
     bench_fn: &str,
     filename: PathBuf,
 ) -> Result<(), String> {
     use inferno::flamegraph::{from_reader, Options};
-    let mut stack = Vec::new();
-    let mut prefix = Vec::new();
-    let mut result = Vec::new();
-    let mut prev = None;
-    for (id, count) in input.into_iter() {
-        if id >= 0 {
-            stack.push((id, count, 0));
-            let name = if id < i32::MAX {
-                match names.get(&id) {
-                    Some(name) => name.clone(),
-                    None => "func_".to_string() + &id.to_string(),
-                }
-            } else {
-                bench_fn.to_string()
-            };
-            prefix.push(name);
-        } else {
-            let end_id = reverse_func_id(id);
-            match stack.pop() {
-                None => return Err("pop empty stack".to_string()),
-                Some((start_id, start, children)) => {
-                    if start_id != end_id {
-                        return Err("func id mismatch".to_string());
-                    }
-                    let cost = count - start;
-                    let frame = prefix.join(";");
-                    prefix.pop().unwrap();
-                    if let Some((parent, parent_cost, children_cost)) = stack.pop() {
-                        stack.push((parent, parent_cost, children_cost + cost));
-                    }
-                    match prev {
-                        Some(prev) if prev == frame => {
-                            // Add an empty spacer to avoid collapsing adjacent same-named calls
-                            // See https://github.com/jonhoo/inferno/issues/185#issuecomment-671393504
-                            result.push(format!("{};spacer 0", prefix.join(";")));
-                        }
-                        _ => (),
-                    }
-                    result.push(format!("{} {}", frame, cost - children));
-                    prev = Some(frame);
-                }
-            }
-        }
-    }
-    let is_trace_incomplete = !stack.is_empty();
+    let mut result = convert_trace_graph_to_flamegraph(input, names, bench_fn);
     let mut opt = Options::default();
     opt.count_name = "instructions".to_string();
-    let bench_fn = if is_trace_incomplete {
-        bench_fn.to_string() + " (incomplete)"
-    } else {
-        bench_fn.to_string()
-    };
-    opt.title = bench_fn;
+    opt.title = bench_fn.to_string();
     opt.flame_chart = true;
     opt.no_sort = true;
     // Reserve result order to make flamegraph from left to right.
@@ -379,6 +330,49 @@ pub(super) fn write_traces_to_file(
     Ok(())
 }
 
+fn convert_trace_graph_to_flamegraph(
+    input: InstructionTraceGraphNode,
+    names: &BTreeMap<i32, String>,
+    bench_fn: &str,
+) -> Vec<String> {
+    let mut flamegraph = Vec::new();
+    let mut prefix = Vec::new();
+    convert_trace_node_to_flamegraph(input, names, bench_fn, &mut prefix, &mut flamegraph);
+    flamegraph
+}
+
+fn convert_trace_node_to_flamegraph(
+    input: InstructionTraceGraphNode,
+    names: &BTreeMap<i32, String>,
+    bench_fn: &str,
+    prefix: &mut Vec<String>,
+    flamegraph: &mut Vec<String>,
+) {
+    let InstructionTraceGraphNode {
+        func_id,
+        cost,
+        children,
+    } = input;
+    prefix.push(func_id_to_name(func_id, names, bench_fn));
+
+    for child in children {
+        convert_trace_node_to_flamegraph(child, names, bench_fn, prefix, flamegraph);
+    }
+    flamegraph.push(format!("{} {}", prefix.join(";"), cost));
+    prefix.pop();
+}
+
+fn func_id_to_name(func_id: i32, names: &BTreeMap<i32, String>, bench_fn: &str) -> String {
+    if func_id < i32::MAX {
+        match names.get(&func_id) {
+            Some(name) => name.clone(),
+            None => "func_".to_string() + &func_id.to_string(),
+        }
+    } else {
+        bench_fn.to_string()
+    }
+}
+
 /// Extracts function names from the module to be a map from function id to function name.
 fn extract_function_names(module: &Module) -> BTreeMap<i32, String> {
     module

canbench-bin/src/lib.rs

@@ -8,7 +8,7 @@ mod results_file;
 mod summary;
 mod table;
 
-use canbench_rs::{BenchResult, Measurement};
+use canbench_rs::{BenchResult, InstructionTraceGraphNode, Measurement};
 use candid::{Encode, Principal};
 use flate2::read::GzDecoder;
 use instruction_tracing::{prepare_instruction_tracing, write_traces_to_file};
@@ -117,14 +117,17 @@ pub fn run_benchmarks(
         }
 
         if let Some(instruction_tracing_canister_id) = instruction_tracing_canister_id {
-            run_instruction_tracing(
-                &pocket_ic,
-                instruction_tracing_canister_id,
-                bench_fn,
-                function_names_mapping.as_ref().unwrap(),
-                results_file,
-                result.total.instructions,
-            );
+            for aggregate in [true, false] {
+                run_instruction_tracing(
+                    &pocket_ic,
+                    instruction_tracing_canister_id,
+                    bench_fn,
+                    function_names_mapping.as_ref().unwrap(),
+                    results_file,
+                    result.total.instructions,
+                    aggregate,
+                );
+            }
         }
 
         new_results.insert(bench_fn.to_string(), result);
@@ -270,15 +273,16 @@ fn run_instruction_tracing(
     names_mapping: &BTreeMap<i32, String>,
     results_file: &Path,
     bench_instructions: u64,
+    aggregate: bool,
 ) {
-    let traces: Result<Vec<(i32, i64)>, String> = match pocket_ic.query_call(
+    let traces: Result<InstructionTraceGraphNode, String> = match pocket_ic.query_call(
         canister_id,
         Principal::anonymous(),
         &format!("__tracing__{bench_fn}"),
-        Encode!(&bench_instructions).unwrap(),
+        Encode!(&bench_instructions, &aggregate).unwrap(),
     ) {
         Ok(reply) => {
-            let res: Result<Vec<(i32, i64)>, String> =
+            let res: Result<InstructionTraceGraphNode, String> =
                 candid::decode_one(&reply).expect("error decoding tracing result");
             res
         }
@@ -295,7 +299,10 @@ fn run_instruction_tracing(
             traces,
             names_mapping,
             bench_fn,
-            results_file.with_file_name(format!("{bench_fn}.svg")),
+            results_file.with_file_name(format!(
+                "{bench_fn}{}.svg",
+                if aggregate { "_aggregated" } else { "" }
+            )),
         )
         .expect("failed to write tracing results"),
         Err(e) => {

canbench-rs-macros/src/lib.rs

@@ -63,9 +63,9 @@ pub fn bench(arg_tokens: TokenStream, item: TokenStream) -> TokenStream {
 
                 #[ic_cdk::query]
                 #[allow(non_snake_case)]
-                fn #tracing_func_name(bench_instructions: u64) -> Result<Vec<(i32, i64)>, String> {
+                fn #tracing_func_name(bench_instructions: u64, aggregate: bool) -> Result<canbench_rs::InstructionTraceGraphNode, String> {
                     #func_name();
-                    canbench_rs::get_traces(bench_instructions)
+                    canbench_rs::get_traces(bench_instructions, aggregate)
                 }
             }
         }
@@ -91,11 +91,11 @@ pub fn bench(arg_tokens: TokenStream, item: TokenStream) -> TokenStream {
 
                 #[ic_cdk::query]
                 #[allow(non_snake_case)]
-                fn #tracing_func_name(bench_instructions: u64) -> Result<Vec<(i32, i64)>, String> {
+                fn #tracing_func_name(bench_instructions: u64, aggregate: bool) -> Result<canbench_rs::InstructionTraceGraphNode, String> {
                     canbench_rs::bench_fn(|| {
                         #func_name();
                     });
-                    canbench_rs::get_traces(bench_instructions)
+                    canbench_rs::get_traces(bench_instructions, aggregate)
                 }
             }
         }