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..96ddce46 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,65 @@ 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 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, 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..b6888ceb 100644 --- a/canbench-bin/tests/tests.rs +++ b/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'" + ); + }); +} diff --git a/canbench-rs/src/lib.rs b/canbench-rs/src/lib.rs index bdc7c20d..ebdd348f 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 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. 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..03d51d44 --- /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 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() {} 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() {}