From 1c9a57a514225aec96ff6e0a673933cfac488772 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Wed, 13 May 2026 16:04:13 +0300 Subject: [PATCH 1/5] feat: Support loading environment variables from a file --- Cargo.lock | 1 + canbench-bin/Cargo.toml | 1 + canbench-bin/src/lib.rs | 74 +++++++++++++++++-- canbench-bin/src/main.rs | 12 +++ canbench-bin/tests/tests.rs | 31 ++++++++ canbench-rs/src/lib.rs | 14 ++++ tests/Cargo.toml | 12 +++ tests/environment_variables/canbench.yml | 6 ++ .../environment_variables.csv | 3 + tests/environment_variables/src/main.rs | 28 +++++++ .../canbench.yml | 6 ++ .../src/main.rs | 1 + .../canbench.yml | 6 ++ .../environment_variables_invalid.csv | 1 + .../environment_variables_invalid/src/main.rs | 1 + 15 files changed, 192 insertions(+), 5 deletions(-) create mode 100644 tests/environment_variables/canbench.yml create mode 100644 tests/environment_variables/environment_variables.csv create mode 100644 tests/environment_variables/src/main.rs create mode 100644 tests/environment_variables_file_does_not_exist/canbench.yml create mode 100644 tests/environment_variables_file_does_not_exist/src/main.rs create mode 100644 tests/environment_variables_invalid/canbench.yml create mode 100644 tests/environment_variables_invalid/environment_variables_invalid.csv create mode 100644 tests/environment_variables_invalid/src/main.rs diff --git a/Cargo.lock b/Cargo.lock index 688b1978..fbc8c279 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -225,6 +225,7 @@ dependencies = [ "colored", "flate2", "hex", + "ic-management-canister-types", "inferno", "pocket-ic", "pretty_assertions", diff --git a/canbench-bin/Cargo.toml b/canbench-bin/Cargo.toml index 687fe244..c2abf944 100644 --- a/canbench-bin/Cargo.toml +++ b/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 diff --git a/canbench-bin/src/lib.rs b/canbench-bin/src/lib.rs index 4b94e5b6..191138c9 100644 --- a/canbench-bin/src/lib.rs +++ b/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, + env_vars_path: Option, 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, instruction_tracing_wasm: Option>, stable_memory_path: Option, + env_vars_path: Option, init_args: Vec, show_canister_output: bool, ) -> (PocketIc, Principal, Option) { @@ -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,58 @@ fn init_pocket_ic( ) } +fn parse_env_vars(env_vars_path: Option) -> Option> { + let env_vars = env_vars_path.map(|path| match std::fs::read(&path) { + Ok(bytes) => { + let mut env_vars = BTreeMap::new(); + for line in String::from_utf8(bytes) + .expect("Environment variables file must be valid UTF-8") + .lines() + { + if let Some((key, value)) = line.split_once(',') { + env_vars.insert(key.trim().to_string(), 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.map(|vars| { + vars.into_iter() + .map(|(name, value)| EnvironmentVariable { name, value }) + .collect() + }) +} + fn init_canister( pocket_ic: &PocketIc, wasm: Vec, init_args: Vec, stable_memory: Option>, + environment_variables: Option>, ) -> 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. diff --git a/canbench-bin/src/main.rs b/canbench-bin/src/main.rs index 07af4283..b28ca623 100644 --- a/canbench-bin/src/main.rs +++ b/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, + + // If provided, the environment variables to set for the canister. + env_vars: Option, } // 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, ); } diff --git a/canbench-bin/tests/tests.rs b/canbench-bin/tests/tests.rs index 471394cd..49dcba35 100644 --- a/canbench-bin/tests/tests.rs +++ b/canbench-bin/tests/tests.rs @@ -390,3 +390,34 @@ 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'" + ); + }); +} diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs index bdc7c20d..b66591be 100644 --- a/canbench-rs/src/lib.rs +++ b/canbench-rs/src/lib.rs @@ -56,6 +56,20 @@ //!
Contents of the stable memory file are loaded after the call to the canister's init method. //! Therefore, changes made to stable memory in the init method would be overwritten.
//! +//! +//! #### Environment Variables +//! +//! A file can be specified to load environment variables from and 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. +//! +//! ```yml +//! environment_variables: +//! 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. diff --git a/tests/Cargo.toml b/tests/Cargo.toml index b22750c5..41b1d674 100644 --- a/tests/Cargo.toml +++ b/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 diff --git a/tests/environment_variables/canbench.yml b/tests/environment_variables/canbench.yml new file mode 100644 index 00000000..df61f379 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/environment_variables/environment_variables.csv b/tests/environment_variables/environment_variables.csv new file mode 100644 index 00000000..93fb9ac2 --- /dev/null +++ b/tests/environment_variables/environment_variables.csv @@ -0,0 +1,3 @@ +name1,value1 with_more_text +name2, value2 +name3 ,value3 \ No newline at end of file diff --git a/tests/environment_variables/src/main.rs b/tests/environment_variables/src/main.rs new file mode 100644 index 00000000..536e4e58 --- /dev/null +++ b/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> = const { RefCell::new(Vec::new()) }; +} + +// A benchmark that prints the environment variables. +#[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() {} diff --git a/tests/environment_variables_file_does_not_exist/canbench.yml b/tests/environment_variables_file_does_not_exist/canbench.yml new file mode 100644 index 00000000..9e3669f1 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/environment_variables_file_does_not_exist/src/main.rs b/tests/environment_variables_file_does_not_exist/src/main.rs new file mode 100644 index 00000000..f328e4d9 --- /dev/null +++ b/tests/environment_variables_file_does_not_exist/src/main.rs @@ -0,0 +1 @@ +fn main() {} diff --git a/tests/environment_variables_invalid/canbench.yml b/tests/environment_variables_invalid/canbench.yml new file mode 100644 index 00000000..ecce47bd --- /dev/null +++ b/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 \ No newline at end of file diff --git a/tests/environment_variables_invalid/environment_variables_invalid.csv b/tests/environment_variables_invalid/environment_variables_invalid.csv new file mode 100644 index 00000000..934edc81 --- /dev/null +++ b/tests/environment_variables_invalid/environment_variables_invalid.csv @@ -0,0 +1 @@ +name \ No newline at end of file diff --git a/tests/environment_variables_invalid/src/main.rs b/tests/environment_variables_invalid/src/main.rs new file mode 100644 index 00000000..f328e4d9 --- /dev/null +++ b/tests/environment_variables_invalid/src/main.rs @@ -0,0 +1 @@ +fn main() {} From e5db1a373492963b961331045f53c22f4c2d7d55 Mon Sep 17 00:00:00 2001 From: mraszyk <31483726+mraszyk@users.noreply.github.com> Date: Wed, 13 May 2026 16:32:04 +0200 Subject: [PATCH 2/5] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- canbench-rs/src/lib.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs index b66591be..f3443b3f 100644 --- a/canbench-rs/src/lib.rs +++ b/canbench-rs/src/lib.rs @@ -64,9 +64,8 @@ //! variable, and `value` is the value of the environment variable. //! //! ```yml -//! environment_variables: -//! file: -//! environment_variables.csv +//! env_vars: +//! file: environment_variables.csv //! ``` //! //! From 3ae19773527956eca20a4ae1538a6ffa91dda3dd Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Wed, 13 May 2026 20:36:32 +0300 Subject: [PATCH 3/5] Apply suggestions from code review Co-authored-by: mraszyk <31483726+mraszyk@users.noreply.github.com> --- canbench-rs/src/lib.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs index f3443b3f..ebdd348f 100644 --- a/canbench-rs/src/lib.rs +++ b/canbench-rs/src/lib.rs @@ -59,9 +59,10 @@ //! //! #### Environment Variables //! -//! A file can be specified to load environment variables from and into the canister. The file +//! 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. +//! variable, and `value` is the value of the environment variable. +//! Leading and trailing whitespaces in `name` and `value` are ignored. //! //! ```yml //! env_vars: From becfcc0443724a31de3d749c4c78e534caed3644 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Wed, 13 May 2026 20:38:08 +0300 Subject: [PATCH 4/5] Update canbench-bin/tests/tests.rs Co-authored-by: mraszyk <31483726+mraszyk@users.noreply.github.com> --- canbench-bin/tests/tests.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/canbench-bin/tests/tests.rs b/canbench-bin/tests/tests.rs index 49dcba35..b6888ceb 100644 --- a/canbench-bin/tests/tests.rs +++ b/canbench-bin/tests/tests.rs @@ -405,8 +405,7 @@ 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 reading environment variables file environment_variables_file_does_not_exist.csv Error: No such file or directory" ); }); From 80f615bb65f4231789875f41970e3f7bbc5a93f8 Mon Sep 17 00:00:00 2001 From: Dimitris Sarlis Date: Wed, 13 May 2026 20:45:07 +0300 Subject: [PATCH 5/5] Review comments --- canbench-bin/src/lib.rs | 29 +++++++++++++++---------- tests/environment_variables/src/main.rs | 2 +- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/canbench-bin/src/lib.rs b/canbench-bin/src/lib.rs index 191138c9..96ddce46 100644 --- a/canbench-bin/src/lib.rs +++ b/canbench-bin/src/lib.rs @@ -432,13 +432,24 @@ fn init_pocket_ic( fn parse_env_vars(env_vars_path: Option) -> Option> { let env_vars = env_vars_path.map(|path| match std::fs::read(&path) { Ok(bytes) => { - let mut env_vars = BTreeMap::new(); - for line in String::from_utf8(bytes) - .expect("Environment variables file must be valid UTF-8") - .lines() - { + 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.insert(key.trim().to_string(), value.trim().to_string()); + env_vars.push(EnvironmentVariable { + name: key.trim().to_string(), + value: value.trim().to_string(), + }); } else { eprintln!( "Invalid line in environment variables file {}: '{}'", @@ -460,11 +471,7 @@ fn parse_env_vars(env_vars_path: Option) -> Option> = const { RefCell::new(Vec::new()) }; } -// A benchmark that prints the environment variables. +// A benchmark that verifies the expected environment variable values were loaded. #[bench] fn state_check() { let state = STATE.with(|s| s.borrow().clone());