Merge branch 'abk/run-522-stable-btreemap-benchmark-in-embedders' int…

…o 'master'

RUN-522: Stable btreemap benchmark

Introduce a benchmark that tests the execution time of an update message which is mutating keys in a `StableBTreeMap`. This benchmark will be expanded in follow-up MRs and used to compare the performance of stable memory with and without Wasm-native stable memory. 

See merge request dfinity-lab/public/ic!10298

.gitlab/CODEOWNERS

@@ -187,6 +187,7 @@ go_deps.bzl               @dfinity-lab/teams/idx
 /rs/rust_canisters/call_tree_test/      @dfinity-lab/teams/execution-owners @dfinity-lab/teams/runtime-owners
 /rs/rust_canisters/proxy_canister/      @dfinity-lab/teams/networking-team
 /rs/rust_canisters/response_payload_test/       @dfinity-lab/teams/runtime-owners
+/rs/rust_canisters/stable_structures/	@dfinity-lab/teams/runtime-owners
 /rs/rust_canisters/xnet_test/           @dfinity-lab/teams/message-routing-owners
 /rs/scenario_tests/                     @dfinity-lab/teams/ic-testing-verification
 /rs/sns/                                @dfinity-lab/teams/nns-team

rs/Cargo.toml

@@ -209,6 +209,7 @@ members = [
   "rosetta-api/icrc1/archive",
   "rosetta-api/hardware_wallet_tests",
   "rosetta-api/test_utils",
+  "rust_canisters/call_tree_test",
   "rust_canisters/canister_creator",
   "rust_canisters/canister_log",
   "rust_canisters/canister_test",
@@ -222,15 +223,15 @@ members = [
   "rust_canisters/ecdsa",
   "rust_canisters/http_types",
   "rust_canisters/memory_test",
-  "rust_canisters/proxy_canister",
   "rust_canisters/on_wire",
   "rust_canisters/pmap",
+  "rust_canisters/proxy_canister",
+  "rust_canisters/response_payload_test",
   "rust_canisters/stable_reader",
+  "rust_canisters/stable_structures",
   "rust_canisters/statesync_test",
   "rust_canisters/tests",
   "rust_canisters/xnet_test",
-  "rust_canisters/call_tree_test",
-  "rust_canisters/response_payload_test",
   "scenario_tests",
   "sns/cli",
   "sns/governance",

rs/embedders/BUILD.bazel

@@ -56,6 +56,18 @@ MACRO_DEV_DEPENDENCIES = []
 
 ALIASES = {}
 
+DATA = [
+    "//rs/canister_sandbox",
+    "//rs/canister_sandbox/sandbox_launcher",
+    "//rs/rust_canisters/stable_structures:stable_structures_canister",
+]
+
+ENV = {
+    "LAUNCHER_BINARY": "$(rootpath //rs/canister_sandbox/sandbox_launcher)",
+    "SANDBOX_BINARY": "$(rootpath //rs/canister_sandbox)",
+    "STABLE_STRUCTURES_CANISTER_WASM_PATH": "$(rootpath //rs/rust_canisters/stable_structures:stable_structures_canister)",
+}
+
 rust_library(
     name = "embedders",
     srcs = glob(["src/**/*.rs"]),
@@ -124,8 +136,21 @@ rust_test_suite_with_extra_srcs(
 )
 
 rust_bench(
-    name = "embedders_bench",
+    name = "compilation_bench",
     srcs = ["benches/compilation.rs"],
     compile_data = glob(["benches/test-data/*"]),
     deps = [":embedders"] + DEPENDENCIES + DEV_DEPENDENCIES,
 )
+
+rust_bench(
+    name = "stable_memory_bench",
+    srcs = ["benches/stable_memory.rs"],
+    data = DATA,
+    env = ENV,
+    deps = [
+        "//rs/rust_canisters/canister_test",
+        "//rs/types/base_types",
+        "//rs/state_machine_tests",
+        "@crate_index//:candid",
+    ] + DEPENDENCIES + DEV_DEPENDENCIES,
+)

rs/embedders/Cargo.toml

@@ -35,8 +35,11 @@ rayon = "1.5.1"
 
 
 [dev-dependencies]
+candid = "0.8.1"
+canister-test = { path = "../rust_canisters/canister_test" }
 criterion = { version = "0.3", features = ["html_reports"] }
 ic-registry-routing-table = { path = "../registry/routing_table" }
+ic-state-machine-tests = { path = "../state_machine_tests" }
 ic-system-api = { path = "../system_api" }
 ic-test-utilities = { path = "../test_utilities" }
 ic-test-utilities-logger = { path = "../test_utilities/logger" }

rs/embedders/benches/stable_memory.rs

@@ -0,0 +1,37 @@
+use candid::Encode;
+use criterion::Criterion;
+
+use ic_state_machine_tests::{Cycles, StateMachine};
+use ic_types::ingress::WasmResult;
+
+const INITIAL_NUMBER_OF_ENTRIES: u64 = 100_000;
+const UPDATE_COUNT: u64 = 10_000;
+
+fn main() {
+    let env = StateMachine::new();
+    let wasm = canister_test::Project::cargo_bin_maybe_from_env("stable_structures_canister", &[]);
+
+    let canister_id = env
+        .install_canister_with_cycles(
+            wasm.bytes(),
+            Encode!(&INITIAL_NUMBER_OF_ENTRIES).unwrap(),
+            None,
+            Cycles::new(1 << 64),
+        )
+        .unwrap();
+
+    let mut criterion = Criterion::default().sample_size(10);
+    let mut group = criterion.benchmark_group("btree");
+    group.bench_function("update", |bench| {
+        bench.iter(|| {
+            let result = env
+                .execute_ingress(
+                    canister_id,
+                    "update_increment_values",
+                    Encode!(&UPDATE_COUNT).unwrap(),
+                )
+                .unwrap();
+            assert!(matches!(result, WasmResult::Reply(_)))
+        });
+    });
+}

rs/rust_canisters/stable_structures/BUILD.bazel

@@ -0,0 +1,15 @@
+load("//bazel:canisters.bzl", "rust_canister")
+
+package(default_visibility = ["//visibility:public"])
+
+rust_canister(
+    name = "stable_structures_canister",
+    srcs = ["src/main.rs"],
+    proc_macro_deps = ["@crate_index//:ic-cdk-macros"],
+    service_file = ":stable_structures.did",
+    deps = [
+        "@crate_index//:candid",
+        "@crate_index//:ic-cdk",
+        "@crate_index//:ic-stable-structures",
+    ],
+)

rs/rust_canisters/stable_structures/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "stable_structures"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+candid = "0.8.1"
+ic-cdk = "0.6.0"
+ic-cdk-macros = "0.6.0"
+ic-stable-structures = "0.5.0"
+serde = "1.0.136"
+
+[[bin]]
+name = "stable_structures"
+path = "src/main.rs"
+

rs/rust_canisters/stable_structures/README.md

@@ -0,0 +1,5 @@
+# Stable Structures Canister
+
+The purpose of this canister is to provide simple examples of uses of
+`ic-stable-structures` in order to benchmark the performance of stable memory
+under realistic workloads.

rs/rust_canisters/stable_structures/src/main.rs

@@ -0,0 +1,34 @@
+use std::{cell::RefCell, ops::Range};
+
+use ic_stable_structures::{DefaultMemoryImpl, StableBTreeMap};
+
+thread_local! {
+    pub static STABLE_MAP: RefCell<StableBTreeMap<u64, u64, DefaultMemoryImpl>> = RefCell::new(
+        StableBTreeMap::init( DefaultMemoryImpl::default())
+    );
+}
+
+fn increment_values(range: Range<u64>) {
+    STABLE_MAP.with(|map| {
+        let mut map = map.borrow_mut();
+        for i in range {
+            match map.get(&i) {
+                None => map.insert(i, i),
+                Some(current) => map.insert(i, current + 1),
+            }
+            .expect("Error inserting into stable map");
+        }
+    })
+}
+
+#[ic_cdk_macros::update]
+fn update_increment_values(entries_to_update: u64) {
+    increment_values(0..entries_to_update)
+}
+
+#[ic_cdk_macros::init]
+fn init(initial_entry_count: u64) {
+    increment_values(0..initial_entry_count)
+}
+
+fn main() {}

rs/rust_canisters/stable_structures/stable_structures.did

@@ -0,0 +1 @@
+service : {}