feat: support loading stable memory from a file

canbench-bin/src/lib.rs

@@ -1,13 +1,14 @@
 //! A module for running benchmarks.
+mod print_benchmark;
+mod results_file;
 use canbench_rs::BenchResult;
 use candid::Principal;
 use flate2::read::GzDecoder;
+use pocket_ic::common::rest::BlobCompression;
 use pocket_ic::{PocketIc, PocketIcBuilder, WasmResult};
-use std::{collections::BTreeMap, env, fs::File, io::Read, path::PathBuf, process::Command};
-mod print_benchmark;
-mod results_file;
 use print_benchmark::print_benchmark;
 use results_file::VersionError;
+use std::{collections::BTreeMap, env, fs::File, io::Read, path::PathBuf, process::Command};
 use wasmparser::Parser as WasmParser;
 
 // The prefix benchmarks are expected to have in their name.
@@ -29,6 +30,7 @@ pub fn run_benchmarks(
     verbose: bool,
     integrity_check: bool,
     runtime_path: &PathBuf,
+    stable_memory_path: Option<PathBuf>,
 ) {
     maybe_download_pocket_ic(runtime_path, verbose, integrity_check);
 
@@ -47,7 +49,12 @@ pub fn run_benchmarks(
     let benchmark_fns = extract_benchmark_fns(canister_wasm_path);
 
     // Initialize PocketIC
-    let (pocket_ic, canister_id) = init_pocket_ic(runtime_path, canister_wasm_path, init_args);
+    let (pocket_ic, canister_id) = init_pocket_ic(
+        runtime_path,
+        canister_wasm_path,
+        stable_memory_path,
+        init_args,
+    );
 
     // Run the benchmarks
     let mut results = BTreeMap::new();
@@ -246,6 +253,7 @@ fn set_env_var_if_unset(key: &str, target_value: &str) {
 fn init_pocket_ic(
     path: &PathBuf,
     canister_wasm_path: &PathBuf,
+    stable_memory_path: Option<PathBuf>,
     init_args: Vec<u8>,
 ) -> (PocketIc, Principal) {
     // PocketIC is used for running the benchmark.
@@ -264,6 +272,28 @@ fn init_pocket_ic(
         init_args,
         None,
     );
+
+    // Load the canister's stable memory if a stable memory file is specified.
+    if let Some(stable_memory_path) = stable_memory_path {
+        let stable_memory_bytes = match std::fs::read(&stable_memory_path) {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                eprintln!(
+                    "Error reading stable memory file {}",
+                    &stable_memory_path.display()
+                );
+                eprintln!("Error: {}", err);
+                std::process::exit(1);
+            }
+        };
+
+        pocket_ic.set_stable_memory(
+            canister_id,
+            stable_memory_bytes,
+            BlobCompression::NoCompression,
+        );
+    }
+
     (pocket_ic, canister_id)
 }
 

canbench-bin/src/main.rs

@@ -43,6 +43,12 @@ struct InitArgs {
     hex: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct StableMemory {
+    // File path to load stable memory from.
+    file: String,
+}
+
 #[derive(Debug, Deserialize)]
 struct Config {
     // If provided, instructs canbench to build the canister
@@ -57,6 +63,9 @@ struct Config {
 
     // If provided, the init arguments to pass to the canister
     init_args: Option<InitArgs>,
+
+    // The stable memory to load into the canister.
+    stable_memory: Option<StableMemory>,
 }
 
 // Path to the canbench directory where we keep internal data.
@@ -113,6 +122,8 @@ fn main() {
         );
     }
 
+    let stable_memory_path = cfg.stable_memory.map(|sm| PathBuf::from(sm.file));
+
     let init_args = cfg
         .init_args
         .map(|args| hex::decode(args.hex).expect("invalid init_args hex value"))
@@ -128,5 +139,6 @@ fn main() {
         !args.less_verbose,
         !args.no_runtime_integrity_check,
         &args.runtime_path.unwrap_or_else(default_runtime_path),
+        stable_memory_path,
     );
 }

canbench-bin/tests/tests.rs

@@ -387,3 +387,24 @@ Benchmark: write_stable_memory (new)
             );
         });
 }
+
+#[test]
+fn loads_stable_memory_file() {
+    BenchTest::canister("stable_memory").run(|output| {
+        // There are assertions in the code of that canister itself, so
+        // all is needed is to assert that the run succeeded.
+        assert_eq!(output.status.code(), Some(0), "output: {:?}", output);
+    });
+}
+
+#[test]
+fn stable_memory_file_not_exit_prints_error() {
+    BenchTest::canister("stable_memory_invalid").run(|output| {
+        assert_err!(
+            output,
+            "
+Error reading stable memory file stable_memory_does_not_exist.bin
+Error: No such file or directory"
+        );
+    });
+}

canbench-rs/src/lib.rs

@@ -42,6 +42,20 @@
 //! init_args:
 //!   hex: 4449444c0001710568656c6c6f
 //! ```
+//!
+//! #### Stable Memory
+//!
+//! A file can be specified to be loaded in the canister's stable memory _after_ initialization.
+//!
+//! ```yml
+//! stable_memory:
+//!   file:
+//!     stable_memory.bin
+//! ```
+//!
+//! <div class="warning">Contents of the stable memory file are loaded <i>after</i> the call to the canister's init method.
+//! Therefore, changes made to stable memory in the init method would be overwritten.</div>
+//!
 //! ### 4. Start benching! 🏋🏽
 //!
 //! Let's say we have a canister that exposes a `query` computing the fibonacci sequence of a given number.

tests/Cargo.toml

@@ -17,6 +17,14 @@ path = "gzipped_wasm/src/main.rs"
 name = "init_arg"
 path = "init_arg/src/main.rs"
 
+[[bin]]
+name = "stable_memory"
+path = "stable_memory/src/main.rs"
+
+[[bin]]
+name = "stable_memory_invalid"
+path = "stable_memory_invalid/src/main.rs"
+
 [dependencies]
 canbench-rs = { path = "../canbench-rs" }
 candid.workspace = true

tests/stable_memory/canbench.yml

@@ -0,0 +1,9 @@
+build_cmd:
+  cargo build --release --target wasm32-unknown-unknown
+
+wasm_path:
+  ../../target/wasm32-unknown-unknown/release/stable_memory.wasm
+
+stable_memory:
+  file:
+    stable_memory.bin

tests/stable_memory/src/main.rs

@@ -0,0 +1,16 @@
+use canbench_rs::bench;
+
+#[bench]
+fn read_from_stable_memory() {
+    let mut buf = [0; 10];
+    ic_cdk::api::stable::stable_read(0, &mut buf);
+
+    // There should be one page in stable memory.
+    assert_eq!(ic_cdk::api::stable::stable_size(), 1);
+
+    // The `stable_memory.bin` specified in canbench.yml only has the first give bytes set.
+    // The rest should be zero.
+    assert_eq!(&buf, &[0x41, 0x42, 0x43, 0x44, 0x45, 0, 0, 0, 0, 0]);
+}
+
+fn main() {}

tests/stable_memory/stable_memory.bin

@@ -0,0 +1 @@
+ABCDE

tests/stable_memory_invalid/canbench.yml

@@ -0,0 +1,9 @@
+build_cmd:
+  cargo build --release --target wasm32-unknown-unknown
+
+wasm_path:
+  ../../target/wasm32-unknown-unknown/release/stable_memory_invalid.wasm
+
+stable_memory:
+  file:
+    stable_memory_does_not_exist.bin

tests/stable_memory_invalid/src/main.rs

@@ -0,0 +1 @@
+fn main() {}