Tracing only subset of functions

.github/workflows/wasm.yml

@@ -9,8 +9,8 @@ jobs:
   deploy:
     runs-on: ubuntu-latest
     env:
-      DFX_VERSION: 0.10.1
-      IC_REPL_VERSION: 0.1.3
+      DFX_VERSION: 0.13.1
+      IC_REPL_VERSION: 0.3.15
     steps:
       - uses: actions/checkout@v2
       - name: Install

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "ic-wasm"
-version = "0.3.4"
+version = "0.3.5"
 authors = ["DFINITY Stiftung"]
 edition = "2021"
 description = "A library for performing Wasm transformations specific to canisters running on the Internet Computer"

README.md

@@ -58,7 +58,17 @@ Usage: `ic-wasm <input.wasm> -o <output.wasm> resource --remove_cycles_transfer
 
 Instrument canister method to emit execution trace to stable memory. 
 
-Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument`
+Usage: `ic-wasm <input.wasm> -o <output.wasm> instrument --trace-only func1 --trace-only func2`
+
+Instrumented canister has the following additional endpoints:
+
+* `__get_cycles: () -> (int64) query`. Get the current cycle counter.
+* `__get_profiling: () -> (vec { record { int32; int64 }}) query`. Get the execution trace log.
+* `__toggle_tracing: () -> ()`. Disable/enable logging the execution trace.
+* `__toggle_entry: () -> ()`. Disable/enable clearing exection trace for each update call.
+* `icp:public name` metadata. Used to map func_id from execution trace to function name.
+
+When `--trace-only` flag is provided, the counter and trace logging will only happen during the execution of that function, instead of tracing the whole update call. Note that the function itself has to be non-recursive.
 
 Current limitations:
 

src/bin/main.rs

@@ -51,7 +51,10 @@ enum SubCommand {
     /// Remove unused functions and debug info
     Shrink,
     /// Instrument canister method to emit execution trace to stable memory (experimental)
-    Instrument,
+    Instrument {
+        #[clap(short, long)]
+        trace_only: Option<Vec<String>>,
+    },
 }
 
 fn main() -> anyhow::Result<()> {
@@ -64,7 +67,11 @@ fn main() -> anyhow::Result<()> {
             ic_wasm::info::info(&m, &mut stdout)?;
         }
         SubCommand::Shrink => ic_wasm::shrink::shrink(&mut m),
-        SubCommand::Instrument => ic_wasm::instrumentation::instrument(&mut m),
+        SubCommand::Instrument { trace_only } => match trace_only {
+            None => ic_wasm::instrumentation::instrument(&mut m, &[]),
+            Some(vec) => ic_wasm::instrumentation::instrument(&mut m, vec),
+        }
+        .map_err(|e| anyhow::anyhow!("{e}"))?,
         SubCommand::Resource {
             remove_cycles_transfer,
             limit_stable_memory_page,

src/instrumentation.rs

@@ -2,6 +2,7 @@ use walrus::ir::*;
 use walrus::*;
 
 use crate::utils::*;
+use std::collections::HashSet;
 
 struct InjectionPoint {
     position: usize,
@@ -28,7 +29,18 @@ struct Variables {
     dynamic_counter64_func: FunctionId,
 }
 
-pub fn instrument(m: &mut Module) {
+/// When trace_only_funcs is not empty, counting and tracing is only enabled for those listed functions per update call.
+/// TODO: doesn't handle recursive entry functions. Need to create a wrapper for the recursive entry function.
+pub fn instrument(m: &mut Module, trace_only_funcs: &[String]) -> Result<(), String> {
+    let mut trace_only_ids = HashSet::new();
+    for name in trace_only_funcs.iter() {
+        let id = match m.funcs.by_name(name) {
+            Some(id) => id,
+            None => return Err(format!("func \"{name}\" not found")),
+        };
+        trace_only_ids.insert(id);
+    }
+    let is_partial_tracing = !trace_only_ids.is_empty();
     let func_cost = FunctionCost::new(m);
     let total_counter = m
         .globals
@@ -45,8 +57,13 @@ pub fn instrument(m: &mut Module) {
     let is_entry = m
         .globals
         .add_local(ValType::I32, true, InitExpr::Value(Value::I32(0)));
-    let dynamic_counter_func = make_dynamic_counter(m, total_counter);
-    let dynamic_counter64_func = make_dynamic_counter64(m, total_counter);
+    let opt_init = if is_partial_tracing {
+        Some(is_init)
+    } else {
+        None
+    };
+    let dynamic_counter_func = make_dynamic_counter(m, total_counter, &opt_init);
+    let dynamic_counter64_func = make_dynamic_counter64(m, total_counter, &opt_init);
     let vars = Variables {
         total_counter,
         log_size,
@@ -59,7 +76,13 @@ pub fn instrument(m: &mut Module) {
 
     for (id, func) in m.funcs.iter_local_mut() {
         if id != dynamic_counter_func && id != dynamic_counter64_func {
-            inject_metering(func, func.entry_block(), &vars, &func_cost);
+            inject_metering(
+                func,
+                func.entry_block(),
+                &vars,
+                &func_cost,
+                is_partial_tracing,
+            );
         }
     }
     let writer = make_stable_writer(m, &vars);
@@ -70,11 +93,14 @@ pub fn instrument(m: &mut Module) {
             && id != dynamic_counter_func
             && id != dynamic_counter64_func
         {
-            inject_profiling_prints(&m.types, printer, id, func);
+            let is_partial_tracing = trace_only_ids.get(&id).is_some();
+            inject_profiling_prints(&m.types, printer, id, func, is_partial_tracing, &vars);
         }
     }
-    //inject_start(m, vars.is_init);
-    inject_init(m, vars.is_init);
+    if !is_partial_tracing {
+        //inject_start(m, vars.is_init);
+        inject_init(m, vars.is_init);
+    }
     inject_canister_methods(m, &vars);
     let leb = make_leb128_encoder(m);
     make_stable_getter(m, &vars, leb);
@@ -83,13 +109,15 @@ pub fn instrument(m: &mut Module) {
     make_toggle_func(m, "__toggle_entry", vars.is_entry);
     let name = make_name_section(m);
     m.customs.add(name);
+    Ok(())
 }
 
 fn inject_metering(
     func: &mut LocalFunction,
     start: InstrSeqId,
     vars: &Variables,
     func_cost: &FunctionCost,
+    is_partial_tracing: bool,
 ) {
     use InjectionKind::*;
     let mut stack = vec![start];
@@ -182,6 +210,19 @@ fn inject_metering(
                     instrs.extend_from_slice(&[
                         (GlobalGet { global: vars.total_counter }.into(), Default::default()),
                         (Const { value: Value::I64(point.cost) }.into(), Default::default()),
+                    ]);
+                    if is_partial_tracing {
+                        #[rustfmt::skip]
+                        instrs.extend_from_slice(&[
+                            (GlobalGet { global: vars.is_init }.into(), Default::default()),
+                            (Const { value: Value::I32(1) }.into(), Default::default()),
+                            (Binop { op: BinaryOp::I32Xor }.into(), Default::default()),
+                            (Unop { op: UnaryOp::I64ExtendUI32 }.into(), Default::default()),
+                            (Binop { op: BinaryOp::I64Mul }.into(), Default::default()),
+                        ]);
+                    }
+                    #[rustfmt::skip]
+                    instrs.extend_from_slice(&[
                         (Binop { op: BinaryOp::I64Add }.into(), Default::default()),
                         (GlobalSet { global: vars.total_counter }.into(), Default::default()),
                     ]);
@@ -208,6 +249,8 @@ fn inject_profiling_prints(
     printer: FunctionId,
     id: FunctionId,
     func: &mut LocalFunction,
+    is_partial_tracing: bool,
+    vars: &Variables,
 ) {
     // Put the original function body inside a block, so that if the code
     // use br_if/br_table to exit the function, we can still output the exit signal.
@@ -230,6 +273,9 @@ fn inject_profiling_prints(
     *(inner_start.instrs_mut()) = start_instrs;
     let inner_start_id = inner_start.id();
     let mut start_builder = func.builder_mut().func_body();
+    if is_partial_tracing {
+        start_builder.i32_const(0).global_set(vars.is_init);
+    }
     start_builder
         .i32_const(id.index() as i32)
         .call(printer)
@@ -239,7 +285,10 @@ fn inject_profiling_prints(
         // TOOD fix when id == 0
         .i32_const(-(id.index() as i32))
         .call(printer);
-
+    // TODO this only works for non-recursive entry function
+    if is_partial_tracing {
+        start_builder.i32_const(1).global_set(vars.is_init);
+    }
     let mut stack = vec![inner_start_id];
     while let Some(seq_id) = stack.pop() {
         let mut builder = func.builder_mut().instr_seq(seq_id);
@@ -307,26 +356,45 @@ fn inject_profiling_prints(
     }
 }
 
-fn make_dynamic_counter(m: &mut Module, total_counter: GlobalId) -> FunctionId {
+fn make_dynamic_counter(
+    m: &mut Module,
+    total_counter: GlobalId,
+    opt_init: &Option<GlobalId>,
+) -> FunctionId {
     let mut builder = FunctionBuilder::new(&mut m.types, &[ValType::I32], &[ValType::I32]);
     let size = m.locals.add(ValType::I32);
-    builder
-        .func_body()
-        .local_get(size)
-        .unop(UnaryOp::I64ExtendUI32)
+    let mut seq = builder.func_body();
+    seq.local_get(size);
+    if let Some(is_init) = opt_init {
+        seq.global_get(*is_init)
+            .i32_const(1)
+            .binop(BinaryOp::I32Xor)
+            .binop(BinaryOp::I32Mul);
+    }
+    seq.unop(UnaryOp::I64ExtendUI32)
         .global_get(total_counter)
         .binop(BinaryOp::I64Add)
         .global_set(total_counter)
         .local_get(size);
     builder.finish(vec![size], &mut m.funcs)
 }
-fn make_dynamic_counter64(m: &mut Module, total_counter: GlobalId) -> FunctionId {
+fn make_dynamic_counter64(
+    m: &mut Module,
+    total_counter: GlobalId,
+    opt_init: &Option<GlobalId>,
+) -> FunctionId {
     let mut builder = FunctionBuilder::new(&mut m.types, &[ValType::I64], &[ValType::I64]);
     let size = m.locals.add(ValType::I64);
-    builder
-        .func_body()
-        .local_get(size)
-        .global_get(total_counter)
+    let mut seq = builder.func_body();
+    seq.local_get(size);
+    if let Some(is_init) = opt_init {
+        seq.global_get(*is_init)
+            .i32_const(1)
+            .binop(BinaryOp::I32Xor)
+            .unop(UnaryOp::I64ExtendUI32)
+            .binop(BinaryOp::I64Mul);
+    }
+    seq.global_get(total_counter)
         .binop(BinaryOp::I64Add)
         .global_set(total_counter)
         .local_get(size);
@@ -357,7 +425,7 @@ fn make_stable_writer(m: &mut Module, vars: &Variables) -> FunctionId {
             |then| {
                 // TODO: This assumes user code doesn't use stable memory
                 then.global_get(vars.page_size)
-                    .i32_const(32)
+                    .i32_const(30) // cannot use the full 2M, because Candid header takes some extra bytes
                     .binop(BinaryOp::I32GtS) // trace >= 2M
                     .if_else(
                         None,

tests/deployable.ic-repl.sh

@@ -65,25 +65,28 @@ function classes_redirect(wasm) {
   S
 };
 
-let S = motoko(file "ok/motoko-instrument.wasm");
+let S = motoko(file("ok/motoko-instrument.wasm"));
 call S.__get_cycles();
 assert _ == (7199 : int64);
-motoko(file "ok/motoko-shrink.wasm");
-motoko(file "ok/motoko-limit.wasm");
+let S = motoko(file("ok/motoko-gc-instrument.wasm"));
+call S.__get_cycles();
+assert _ == (177 : int64);
+motoko(file("ok/motoko-shrink.wasm"));
+motoko(file("ok/motoko-limit.wasm"));
 
-let S = rust(file "ok/rust-instrument.wasm");
+let S = rust(file("ok/rust-instrument.wasm"));
 call S.__get_cycles();
 assert _ == (66016 : int64);
-rust(file "ok/rust-shrink.wasm");
-rust(file "ok/rust-limit.wasm");
+rust(file("ok/rust-shrink.wasm"));
+rust(file("ok/rust-limit.wasm"));
 
-let S = wat(file "ok/wat-instrument.wasm");
+let S = wat(file("ok/wat-instrument.wasm"));
 call S.__get_cycles();
 assert _ == (121 : int64);
-wat(file "ok/wat-shrink.wasm");
-wat(file "ok/wat-limit.wasm");
+wat(file("ok/wat-shrink.wasm"));
+wat(file("ok/wat-limit.wasm"));
 
-classes(file "ok/classes-shrink.wasm");
-classes_limit(file "ok/classes-limit.wasm");
-classes_redirect(file "ok/classes-redirect.wasm");
-classes(file "ok/classes-nop-redirect.wasm");
+classes(file("ok/classes-shrink.wasm"));
+classes_limit(file("ok/classes-limit.wasm"));
+classes_redirect(file("ok/classes-redirect.wasm"));
+classes(file("ok/classes-nop-redirect.wasm"));

tests/tests.rs

@@ -34,6 +34,13 @@ fn instrumentation() {
         .assert()
         .success();
     assert_wasm("motoko-instrument.wasm");
+    wasm_input("motoko.wasm", true)
+        .arg("instrument")
+        .arg("-t")
+        .arg("schedule_copying_gc")
+        .assert()
+        .success();
+    assert_wasm("motoko-gc-instrument.wasm");
     wasm_input("wat.wasm", true)
         .arg("instrument")
         .assert()