test: Add a few new benchmarks that exercise the MemoryManager which is often how StableStructures are used (#253)

Also fix the 10mb benchmark, previously we added 200 10**kb** values,
now we'll add 20 10**mb** values (like the benchmark description says).
200 10mb values crash the benchmark so let's stick with 20.

Author: frankdavid

benchmarks/src/btreemap.rs

@@ -1,6 +1,7 @@
 use crate::Random;
 use canbench_rs::{bench, bench_fn, BenchResult};
-use ic_stable_structures::{storable::Blob, BTreeMap, DefaultMemoryImpl, Storable};
+use ic_stable_structures::memory_manager::{MemoryId, MemoryManager};
+use ic_stable_structures::{storable::Blob, BTreeMap, DefaultMemoryImpl, Memory, Storable};
 use std::ops::Bound;
 use tiny_rng::{Rand, Rng};
 
@@ -164,12 +165,24 @@ pub fn btreemap_insert_blob_1024_512_v2() -> BenchResult {
     insert_blob_helper_v2::<1024, 512>()
 }
 
+#[bench(raw)]
+pub fn btreemap_insert_blob_1024_512_v2_mem_manager() -> BenchResult {
+    insert_blob_helper_v2_mem_manager::<1024, 512>()
+}
+
 #[bench(raw)]
 pub fn btreemap_insert_u64_u64() -> BenchResult {
     let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());
     insert_helper::<u64, u64>(btree)
 }
 
+#[bench(raw)]
+pub fn btreemap_insert_u64_u64_mem_manager() -> BenchResult {
+    let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    let btree = BTreeMap::new(memory_manager.get(MemoryId::new(42)));
+    insert_helper::<u64, u64>(btree)
+}
+
 #[bench(raw)]
 pub fn btreemap_insert_u64_u64_v2() -> BenchResult {
     let btree = BTreeMap::new(DefaultMemoryImpl::default());
@@ -204,16 +217,20 @@ pub fn btreemap_insert_blob_8_u64_v2() -> BenchResult {
 pub fn btreemap_insert_10mib_values() -> BenchResult {
     let mut btree = BTreeMap::new(DefaultMemoryImpl::default());
 
-    // Insert 200 10MiB values.
+    // Insert 20 10MiB values.
     let mut rng = Rng::from_seed(0);
     let mut values = vec![];
-    for _ in 0..200 {
-        values.push(rng.iter(Rand::rand_u8).take(10 * 1024).collect::<Vec<_>>());
+    for _ in 0..20 {
+        values.push(
+            rng.iter(Rand::rand_u8)
+                .take(10 * 1024 * 1024)
+                .collect::<Vec<_>>(),
+        );
     }
 
     bench_fn(|| {
         let mut i = 0u64;
-        for value in values.into_iter() {
+        for value in values {
             btree.insert(i, value);
             i += 1;
         }
@@ -543,6 +560,11 @@ pub fn btreemap_get_blob_512_1024_v2() -> BenchResult {
     get_blob_helper_v2::<512, 1024>()
 }
 
+#[bench(raw)]
+pub fn btreemap_get_blob_512_1024_v2_mem_manager() -> BenchResult {
+    get_blob_helper_v2_mem_manager::<512, 1024>()
+}
+
 #[bench(raw)]
 pub fn btreemap_get_u64_u64() -> BenchResult {
     let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());
@@ -555,6 +577,13 @@ pub fn btreemap_get_u64_u64_v2() -> BenchResult {
     get_helper::<u64, u64>(btree)
 }
 
+#[bench(raw)]
+pub fn btreemap_get_u64_u64_v2_mem_manager() -> BenchResult {
+    let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    let btree = BTreeMap::new(memory_manager.get(MemoryId::new(42)));
+    get_helper::<u64, u64>(btree)
+}
+
 #[bench(raw)]
 pub fn btreemap_get_u64_blob_8() -> BenchResult {
     let btree = BTreeMap::new_v1(DefaultMemoryImpl::default());
@@ -590,9 +619,15 @@ fn insert_blob_helper_v2<const K: usize, const V: usize>() -> BenchResult {
     insert_helper::<Blob<K>, Blob<V>>(btree)
 }
 
+fn insert_blob_helper_v2_mem_manager<const K: usize, const V: usize>() -> BenchResult {
+    let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    let btree = BTreeMap::new(memory_manager.get(MemoryId::new(42)));
+    insert_helper::<Blob<K>, Blob<V>>(btree)
+}
+
 // Profiles inserting a large number of random blobs into a btreemap.
 fn insert_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
-    mut btree: BTreeMap<K, V, DefaultMemoryImpl>,
+    mut btree: BTreeMap<K, V, impl Memory>,
 ) -> BenchResult {
     let num_keys = 10_000;
     let mut rng = Rng::from_seed(0);
@@ -640,8 +675,14 @@ fn get_blob_helper_v2<const K: usize, const V: usize>() -> BenchResult {
     get_helper::<Blob<K>, Blob<V>>(btree)
 }
 
+fn get_blob_helper_v2_mem_manager<const K: usize, const V: usize>() -> BenchResult {
+    let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    let btree = BTreeMap::new(memory_manager.get(MemoryId::new(42)));
+    get_helper::<Blob<K>, Blob<V>>(btree)
+}
+
 fn get_helper<K: Clone + Ord + Storable + Random, V: Storable + Random>(
-    mut btree: BTreeMap<K, V, DefaultMemoryImpl>,
+    mut btree: BTreeMap<K, V, impl Memory>,
 ) -> BenchResult {
     let num_keys = 10_000;
     let mut rng = Rng::from_seed(0);

benchmarks/src/vec.rs

@@ -1,7 +1,8 @@
 use crate::Random;
 use canbench_rs::{bench, bench_fn, BenchResult};
+use ic_stable_structures::memory_manager::{MemoryId, MemoryManager};
 use ic_stable_structures::storable::Blob;
-use ic_stable_structures::{DefaultMemoryImpl, StableVec, Storable};
+use ic_stable_structures::{DefaultMemoryImpl, Memory, StableVec, Storable};
 use tiny_rng::{Rand, Rng};
 
 #[bench(raw)]
@@ -44,6 +45,11 @@ pub fn vec_get_blob_4() -> BenchResult {
     vec_get_blob::<4>()
 }
 
+#[bench(raw)]
+pub fn vec_get_blob_4_mem_manager() -> BenchResult {
+    vec_get_blob_mem_manager::<4>()
+}
+
 #[bench(raw)]
 pub fn vec_get_blob_8() -> BenchResult {
     vec_get_blob::<8>()
@@ -64,14 +70,19 @@ pub fn vec_get_blob_64() -> BenchResult {
     vec_get_blob::<64>()
 }
 
+#[bench(raw)]
+pub fn vec_get_blob_64_mem_manager() -> BenchResult {
+    vec_get_blob_mem_manager::<64>()
+}
+
 #[bench(raw)]
 pub fn vec_get_blob_128() -> BenchResult {
     vec_get_blob::<128>()
 }
 
 #[bench(raw)]
 pub fn vec_get_u64() -> BenchResult {
-    vec_get::<u64>()
+    vec_get::<u64>(DefaultMemoryImpl::default())
 }
 
 fn vec_insert_blob<const N: usize>() -> BenchResult {
@@ -97,12 +108,17 @@ fn vec_insert<T: Storable + Random>() -> BenchResult {
 }
 
 fn vec_get_blob<const N: usize>() -> BenchResult {
-    vec_get::<Blob<N>>()
+    vec_get::<Blob<N>>(DefaultMemoryImpl::default())
+}
+
+fn vec_get_blob_mem_manager<const N: usize>() -> BenchResult {
+    let memory_manager = MemoryManager::init(DefaultMemoryImpl::default());
+    vec_get::<Blob<N>>(memory_manager.get(MemoryId::new(42)))
 }
 
-fn vec_get<T: Storable + Random>() -> BenchResult {
+fn vec_get<T: Storable + Random>(memory: impl Memory) -> BenchResult {
     let num_items = 10_000;
-    let svec: StableVec<T, _> = StableVec::new(DefaultMemoryImpl::default()).unwrap();
+    let svec: StableVec<T, _> = StableVec::new(memory).unwrap();
 
     let mut rng = Rng::from_seed(0);
 

canbench_results.yml

@@ -71,6 +71,12 @@ benches:
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
+  btreemap_get_blob_512_1024_v2_mem_manager:
+    total:
+      instructions: 3164909208
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: { }
   btreemap_get_blob_64_1024:
     total:
       instructions: 572513751
@@ -130,12 +136,18 @@ benches:
       instructions: 255485892
       heap_increase: 0
       stable_memory_increase: 0
+    scopes: { }
+  btreemap_get_u64_u64_v2_mem_manager:
+    total:
+      instructions: 673752904
+      heap_increase: 0
+      stable_memory_increase: 0
     scopes: {}
   btreemap_insert_10mib_values:
     total:
-      instructions: 82015559
+      instructions: 5251584111
       heap_increase: 0
-      stable_memory_increase: 32
+      stable_memory_increase: 3613
     scopes: {}
   btreemap_insert_blob_1024_128:
     total:
@@ -208,6 +220,12 @@ benches:
       instructions: 5070548871
       heap_increase: 0
       stable_memory_increase: 261
+    scopes: { }
+  btreemap_insert_blob_1024_512_v2_mem_manager:
+    total:
+      instructions: 6292667075
+      heap_increase: 0
+      stable_memory_increase: 256
     scopes: {}
   btreemap_insert_blob_1024_64:
     total:
@@ -359,6 +377,12 @@ benches:
       heap_increase: 0
       stable_memory_increase: 7
     scopes: {}
+  btreemap_insert_u64_u64_mem_manager:
+    total:
+      instructions: 1019612837
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: { }
   btreemap_insert_u64_u64_v2:
     total:
       instructions: 421501977
@@ -635,11 +659,23 @@ benches:
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
+  vec_get_blob_4_mem_manager:
+    total:
+      instructions: 14476122
+      heap_increase: 0
+      stable_memory_increase: 0
+    scopes: { }
   vec_get_blob_64:
     total:
       instructions: 12960294
       heap_increase: 0
       stable_memory_increase: 0
+    scopes: { }
+  vec_get_blob_64_mem_manager:
+    total:
+      instructions: 24146389
+      heap_increase: 0
+      stable_memory_increase: 0
     scopes: {}
   vec_get_blob_8:
     total: