feat: Support loading environment variables from a file

canbench-bin/Cargo.toml

@@ -27,6 +27,7 @@ inferno = { version = "0.11", default-features = false, features = [
 # `pocket-ic` should be pinned to an exact version so that the PocketIC server binary version
 # `POCKET_IC_SERVER_VERSION` defined in `canbench-bin/src/lib.rs` is compatible.
 pocket-ic = "=13.0.0"
+ic-management-canister-types = "0.5.0"
 reqwest.workspace = true
 rustc-demangle.workspace = true
 semver.workspace = true

canbench-bin/src/lib.rs

@@ -11,9 +11,10 @@ mod table;
 use canbench_rs::{BenchResult, Measurement};
 use candid::{Encode, Principal};
 use flate2::read::GzDecoder;
+use ic_management_canister_types::EnvironmentVariable;
 use instruction_tracing::{prepare_instruction_tracing, write_traces_to_file};
 use pocket_ic::common::rest::BlobCompression;
-use pocket_ic::{PocketIc, PocketIcBuilder};
+use pocket_ic::{CanisterSettings, PocketIc, PocketIcBuilder};
 use print_benchmark::print_benchmark;
 use results_file::VersionError;
 use std::{
@@ -56,6 +57,7 @@ pub fn run_benchmarks(
     instruction_tracing: bool,
     runtime_path: &PathBuf,
     stable_memory_path: Option<PathBuf>,
+    env_vars_path: Option<PathBuf>,
     noise_threshold: f64,
 ) {
     maybe_download_pocket_ic(runtime_path, verbose, integrity_check);
@@ -90,6 +92,7 @@ pub fn run_benchmarks(
         benchmark_wasm,
         instruction_tracing_wasm,
         stable_memory_path,
+        env_vars_path,
         init_args,
         show_canister_output,
     );
@@ -376,6 +379,7 @@ fn init_pocket_ic(
     benchmark_wasm: Vec<u8>,
     instruction_tracing_wasm: Option<Vec<u8>>,
     stable_memory_path: Option<PathBuf>,
+    env_vars_path: Option<PathBuf>,
     init_args: Vec<u8>,
     show_canister_output: bool,
 ) -> (PocketIc, Principal, Option<Principal>) {
@@ -399,9 +403,24 @@ fn init_pocket_ic(
         }
     });
 
-    let instruction_tracing_canister_id = instruction_tracing_wasm
-        .map(|wasm| init_canister(&pocket_ic, wasm, init_args.clone(), stable_memory.clone()));
-    let benchmark_canister_id = init_canister(&pocket_ic, benchmark_wasm, init_args, stable_memory);
+    let environment_variables = parse_env_vars(env_vars_path);
+
+    let instruction_tracing_canister_id = instruction_tracing_wasm.map(|wasm| {
+        init_canister(
+            &pocket_ic,
+            wasm,
+            init_args.clone(),
+            stable_memory.clone(),
+            environment_variables.clone(),
+        )
+    });
+    let benchmark_canister_id = init_canister(
+        &pocket_ic,
+        benchmark_wasm,
+        init_args,
+        stable_memory,
+        environment_variables,
+    );
 
     (
         pocket_ic,
@@ -410,13 +429,65 @@ fn init_pocket_ic(
     )
 }
 
+fn parse_env_vars(env_vars_path: Option<PathBuf>) -> Option<Vec<EnvironmentVariable>> {
+    let env_vars = env_vars_path.map(|path| match std::fs::read(&path) {
+        Ok(bytes) => {
+            let contents = match String::from_utf8(bytes) {
+                Ok(contents) => contents,
+                Err(err) => {
+                    eprintln!(
+                        "Environment variables file {} is not valid UTF-8",
+                        path.display()
+                    );
+                    eprintln!("Error: {}", err);
+                    std::process::exit(1);
+                }
+            };
+            let mut env_vars = Vec::new();
+            for line in contents.lines() {
+                if let Some((key, value)) = line.split_once(',') {
+                    env_vars.push(EnvironmentVariable {
+                        name: key.trim().to_string(),
+                        value: value.trim().to_string(),
+                    });
+                } else {
+                    eprintln!(
+                        "Invalid line in environment variables file {}: '{}'",
+                        path.display(),
+                        line
+                    );
+                    std::process::exit(1);
+                }
+            }
+            env_vars
+        }
+        Err(err) => {
+            eprintln!(
+                "Error reading environment variables file {}",
+                path.display()
+            );
+            eprintln!("Error: {}", err);
+            std::process::exit(1);
+        }
+    });
+
+    env_vars
+}
+
 fn init_canister(
     pocket_ic: &PocketIc,
     wasm: Vec<u8>,
     init_args: Vec<u8>,
     stable_memory: Option<Vec<u8>>,
+    environment_variables: Option<Vec<EnvironmentVariable>>,
 ) -> Principal {
-    let canister_id = pocket_ic.create_canister();
+    let canister_id = pocket_ic.create_canister_with_settings(
+        None,
+        Some(CanisterSettings {
+            environment_variables,
+            ..Default::default()
+        }),
+    );
     pocket_ic.add_cycles(canister_id, 1_000_000_000_000_000);
     pocket_ic.install_canister(canister_id, wasm, init_args, None);
     // Load the canister's stable memory if stable memory is specified.

canbench-bin/src/main.rs

@@ -76,6 +76,12 @@ struct StableMemory {
     file: String,
 }
 
+#[derive(Debug, Deserialize)]
+struct EnvironmentVariables {
+    // File path to load environment variables from.
+    file: String,
+}
+
 #[derive(Debug, Deserialize)]
 struct Config {
     // If provided, instructs canbench to build the canister
@@ -97,6 +103,9 @@ struct Config {
 
     // The stable memory to load into the canister.
     stable_memory: Option<StableMemory>,
+
+    // If provided, the environment variables to set for the canister.
+    env_vars: Option<EnvironmentVariables>,
 }
 
 // Path to the canbench directory where we keep internal data.
@@ -168,6 +177,8 @@ fn main() {
         .map(|args| hex::decode(args.hex).expect("invalid init_args hex value"))
         .unwrap_or_default();
 
+    let env_vars_path = cfg.env_vars.map(|ev| PathBuf::from(ev.file));
+
     // Run the benchmarks.
     canbench::run_benchmarks(
         &wasm_path,
@@ -185,6 +196,7 @@ fn main() {
         args.instruction_tracing,
         &args.runtime_path.unwrap_or_else(default_runtime_path),
         stable_memory_path,
+        env_vars_path,
         args.noise_threshold,
     );
 }

canbench-bin/tests/tests.rs

@@ -390,3 +390,33 @@ fn reports_recursive_scopes_benchmark() {
             assert_success!(output, expected.as_str());
         });
 }
+
+#[test]
+fn loads_environment_variables_file() {
+    BenchTest::canister("environment_variables").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 environment_variables_file_does_not_exist_prints_error() {
+    BenchTest::canister("environment_variables_file_does_not_exist").run(|output| {
+        assert_err!(
+            output,
+            "Error reading environment variables file environment_variables_file_does_not_exist.csv
+Error: No such file or directory"
+        );
+    });
+}
+
+#[test]
+fn environment_variables_invalid_file_prints_error() {
+    BenchTest::canister("environment_variables_invalid").run(|output| {
+        assert_err!(
+            output,
+            "Invalid line in environment variables file environment_variables_invalid.csv: 'name'"
+        );
+    });
+}

canbench-rs/src/lib.rs

@@ -56,6 +56,20 @@
 //! <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>
 //!
+//!
+//! #### Environment Variables
+//!
+//! A file can be specified from which environment variables are loaded into the canister. The file
+//! is a CSV with two columns: `name` and `value`, where `name` is the name of the environment
+//! variable, and `value` is the value of the environment variable.
+//! Leading and trailing whitespaces in `name` and `value` are ignored.
+//!
+//! ```yml
+//! env_vars:
+//!   file: environment_variables.csv
+//! ```
+//!
+//!
 //! ### 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

@@ -29,6 +29,18 @@ path = "stable_memory/src/main.rs"
 name = "stable_memory_invalid"
 path = "stable_memory_invalid/src/main.rs"
 
+[[bin]]
+name = "environment_variables"
+path = "environment_variables/src/main.rs"
+
+[[bin]]
+name = "environment_variables_file_does_not_exist"
+path = "environment_variables_file_does_not_exist/src/main.rs"
+
+[[bin]]
+name = "environment_variables_invalid"
+path = "environment_variables_invalid/src/main.rs"
+
 [dependencies]
 canbench-rs = { path = "../canbench-rs" }
 candid.workspace = true

tests/environment_variables/canbench.yml

@@ -0,0 +1,6 @@
+build_cmd: cargo build --release --target wasm32-unknown-unknown --locked
+
+wasm_path: ../../target/wasm32-unknown-unknown/release/environment_variables.wasm
+
+env_vars:
+  file: environment_variables.csv

tests/environment_variables/environment_variables.csv

@@ -0,0 +1,3 @@
+name1,value1 with_more_text
+name2, value2
+name3 ,value3

tests/environment_variables/src/main.rs

@@ -0,0 +1,28 @@
+use canbench_rs::bench;
+use ic_cdk::api::env_var_value;
+use std::cell::RefCell;
+
+thread_local! {
+    static STATE: RefCell<Vec<String>> = const { RefCell::new(Vec::new()) };
+}
+
+// A benchmark that verifies the expected environment variable values were loaded.
+#[bench]
+fn state_check() {
+    let state = STATE.with(|s| s.borrow().clone());
+    assert_eq!(state[0], "value1 with_more_text");
+    assert_eq!(state[1], "value2");
+    assert_eq!(state[2], "value3");
+}
+
+#[ic_cdk::init]
+pub fn init() {
+    STATE.with(|s| {
+        let mut state = s.borrow_mut();
+        state.push(env_var_value("name1"));
+        state.push(env_var_value("name2"));
+        state.push(env_var_value("name3"));
+    });
+}
+
+fn main() {}

tests/environment_variables_file_does_not_exist/canbench.yml

@@ -0,0 +1,6 @@
+build_cmd: cargo build --release --target wasm32-unknown-unknown --locked
+
+wasm_path: ../../target/wasm32-unknown-unknown/release/environment_variables_file_does_not_exist.wasm
+
+env_vars:
+  file: environment_variables_file_does_not_exist.csv

tests/environment_variables_file_does_not_exist/src/main.rs

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

tests/environment_variables_invalid/canbench.yml

@@ -0,0 +1,6 @@
+build_cmd: cargo build --release --target wasm32-unknown-unknown --locked
+
+wasm_path: ../../target/wasm32-unknown-unknown/release/environment_variables_invalid.wasm
+
+env_vars:
+  file: environment_variables_invalid.csv

tests/environment_variables_invalid/environment_variables_invalid.csv

@@ -0,0 +1 @@
+name

tests/environment_variables_invalid/src/main.rs

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