feat: add basic profiling to benchmarks

Cargo.toml

@@ -11,6 +11,10 @@ keywords = ["internet-computer", "dfinity", "stable-structures"]
 include = ["src", "Cargo.toml", "LICENSE", "README.md"]
 repository = "https://github.com/dfinity/stable-structures"
 
+[dependencies]
+# An optional dependency to profile parts of the code.
+profiler = { path = "./profiler", optional = true }
+
 [dev-dependencies]
 criterion = "0.4.0"
 ic-cdk = "0.6.8"

benches/benches.rs

@@ -144,6 +144,10 @@ fn execution_instructions(arguments: ExecutionArguments) -> u64 {
     let stderr = String::from_utf8(output.stderr).unwrap();
     assert!(output.status.success(), "{stdout}\n{stderr}");
 
+    // Output logs from the canister as it contains profiling data.
+    println!();
+    println!("{stderr}");
+
     // Convert result formatted as "(1_000_000 : nat64)" to u64.
     let result = stdout
         .trim()

benchmark-canisters/Cargo.toml

@@ -8,5 +8,6 @@ edition = "2021"
 [dependencies]
 ic-cdk = "0.6.8"
 ic-cdk-macros = "0.6.8"
-ic-stable-structures = { path = "../" }
+ic-stable-structures = { path = "../", features = ["profiler"] }
+profiler = { path = "../profiler" }
 tiny-rng = "0.2.0"

benchmark-canisters/src/main.rs

@@ -8,8 +8,32 @@ mod vec;
 /// Returns the number of instructions consumed by the given function.
 pub(crate) fn count_instructions<R>(f: impl FnOnce() -> R) -> u64 {
     let start = ic_cdk::api::performance_counter(0);
+    profiler::reset();
     f();
-    ic_cdk::api::performance_counter(0) - start
+    let total_instructions = ic_cdk::api::performance_counter(0) - start;
+
+    let profiling_results: std::collections::BTreeMap<_, _> = profiler::get_results()
+        .into_iter()
+        .map(|(k, v)| {
+            (
+                k,
+                format!("{} ({}%)", format_num(v), ((v * 100) / total_instructions)),
+            )
+        })
+        .collect();
+    ic_cdk::api::print(&format!("{:#?}", profiling_results));
+    total_instructions
+}
+
+fn format_num(num: u64) -> String {
+    num.to_string()
+        .as_bytes()
+        .rchunks(3)
+        .rev()
+        .map(std::str::from_utf8)
+        .collect::<Result<Vec<&str>, _>>()
+        .unwrap()
+        .join("_")
 }
 
 const fn max_size<A: Storable>() -> u32 {

profiler/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "profiler"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+ic-cdk = "0.6.8"

profiler/src/lib.rs

@@ -0,0 +1,64 @@
+//! A module for profiling canisters.
+use std::cell::RefCell;
+use std::collections::BTreeMap;
+
+thread_local! {
+    static PROFILING: RefCell<BTreeMap<&'static str, u64>> = RefCell::new(BTreeMap::new());
+}
+
+/// Starts profiling the instructions consumed.
+///
+/// Instructions are counted and recorded under the given name until the
+/// `Profile` object returned is dropped.
+pub fn profile(name: &'static str) -> Profile {
+    Profile::new(name)
+}
+
+/// Clears all profiling data.
+pub fn reset() {
+    PROFILING.with(|p| p.borrow_mut().clear());
+}
+
+/// Returns the number of instructions used for each of the profile names.
+pub fn get_results() -> std::collections::BTreeMap<&'static str, u64> {
+    PROFILING.with(|p| p.borrow().clone())
+}
+
+pub struct Profile {
+    name: &'static str,
+    start_instructions: u64,
+}
+
+impl Profile {
+    fn new(name: &'static str) -> Self {
+        Self {
+            name,
+            start_instructions: instruction_count(),
+        }
+    }
+}
+
+impl Drop for Profile {
+    fn drop(&mut self) {
+        let instructions_count = instruction_count() - self.start_instructions;
+
+        PROFILING.with(|p| {
+            let mut p = p.borrow_mut();
+            let entry = p.entry(&self.name).or_insert(0);
+            *entry += instructions_count;
+        });
+    }
+}
+
+fn instruction_count() -> u64 {
+    #[cfg(target_arch = "wasm32")]
+    {
+        ic_cdk::api::performance_counter(0)
+    }
+
+    #[cfg(not(target_arch = "wasm32"))]
+    {
+        // Consider using cpu time here.
+        0
+    }
+}

src/btreemap.rs

@@ -1309,8 +1309,6 @@ mod test {
 
             // Overwrite the value of the median key.
             assert_eq!(btree.insert(b(&[12]), b(&[1, 2, 3])), Some(b(&[])));
-
-            // The value is overwritten successfully.
             assert_eq!(btree.get(&b(&[12])), Some(b(&[1, 2, 3])));
 
             // The child has not been split and is still full.
@@ -2157,7 +2155,6 @@ mod test {
                 vec![(b(&[1, 4]), vec![]), (b(&[2, 3]), vec![])]
             );
             assert_eq!(root.children_len(), 3);
-
             let child_0 = btree.load_node(root.child(0));
             assert_eq!(child_0.node_type(), NodeType::Leaf);
             assert_eq!(

src/btreemap/node/v1.rs

@@ -59,6 +59,9 @@ impl<K: Storable + Ord + Clone> Node<K> {
         max_value_size: u32,
         memory: &M,
     ) -> Self {
+        #[cfg(feature = "profiler")]
+        let _p = profiler::profile("node_load_v1");
+
         // Load the header.
         let header: NodeHeader = read_struct(address, memory);
         assert_eq!(&header.magic, MAGIC, "Bad magic.");
@@ -118,6 +121,9 @@ impl<K: Storable + Ord + Clone> Node<K> {
     }
 
     pub(super) fn save_v1<M: Memory>(&self, memory: &M) {
+        #[cfg(feature = "profiler")]
+        let _p = profiler::profile("node_save_v1");
+
         match self.node_type {
             NodeType::Leaf => {
                 assert!(self.children.is_empty());

src/btreemap/node/v2.rs

@@ -112,6 +112,9 @@ impl<K: Storable + Ord + Clone> Node<K> {
 
     /// Loads a v2 node from memory at the given address.
     pub(super) fn load_v2<M: Memory>(address: Address, page_size: PageSize, memory: &M) -> Self {
+        #[cfg(feature = "profiler")]
+        let _p = profiler::profile("node_load_v2");
+
         // Load the node, including any overflows, into a buffer.
         let node_buf = read_node(address, page_size.get(), memory);
 
@@ -192,6 +195,9 @@ impl<K: Storable + Ord + Clone> Node<K> {
 
     // Saves the node to memory.
     pub(super) fn save_v2<M: Memory>(&self, allocator: &mut Allocator<M>) {
+        #[cfg(feature = "profiler")]
+        let _p = profiler::profile("node_save_v2");
+
         let page_size = self.version.page_size().get();
         assert!(page_size >= MINIMUM_PAGE_SIZE);
         assert_eq!(self.keys.len(), self.encoded_values.borrow().len());
@@ -255,7 +261,7 @@ impl<K: Storable + Ord + Clone> Node<K> {
     ) {
         // Compute how many overflow pages are needed.
         let additional_pages_needed = if buf.len() > page_size {
-            debug_assert!(page_size >= PAGE_OVERFLOW_DATA_OFFSET.into());
+            debug_assert!(page_size >= PAGE_OVERFLOW_DATA_OFFSET.get() as usize);
             let overflow_page_capacity = page_size - PAGE_OVERFLOW_DATA_OFFSET.get() as usize;
             let overflow_data_len = buf.len() - page_size;