sparse-merkle-tree bench update (aptos-labs#8613)

* jemalloc for smt bench * update * tweak
zi0Black · Jun 13, 2023 · 3008560 · 3008560
1 parent 3b05c8f
commit 3008560
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 16 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/storage/scratchpad/Cargo.toml b/storage/scratchpad/Cargo.toml
@@ -20,6 +20,7 @@ aptos-types = { workspace = true }
 bitvec = { workspace = true }
 criterion = { workspace = true, optional = true }
 itertools = { workspace = true }
+jemallocator = { workspace = true }
 once_cell = { workspace = true }
 proptest = { workspace = true, optional = true }
 rayon = { workspace = true }

diff --git a/storage/scratchpad/benches/sparse_merkle.rs b/storage/scratchpad/benches/sparse_merkle.rs
@@ -13,22 +13,19 @@ use itertools::zip_eq;
 use rand::{distributions::Standard, prelude::StdRng, seq::IteratorRandom, Rng, SeedableRng};
 use std::collections::HashSet;
 
+#[cfg(unix)]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
 struct Block {
     smt: SparseMerkleTree<StateValue>,
-    updates: Vec<Vec<(HashValue, Option<StateValue>)>>,
+    updates: Vec<(HashValue, Option<StateValue>)>,
     proof_reader: ProofReader,
 }
 
 impl Block {
-    fn updates(&self) -> Vec<Vec<(HashValue, Option<&StateValue>)>> {
-        self.updates
-            .iter()
-            .map(|small_batch| small_batch.iter().map(|(k, v)| (*k, v.as_ref())).collect())
-            .collect()
-    }
-
-    fn updates_flat_batch(&self) -> Vec<(HashValue, Option<&StateValue>)> {
-        self.updates().iter().flatten().cloned().collect()
+    fn updates(&self) -> Vec<(HashValue, Option<&StateValue>)> {
+        self.updates.iter().map(|(k, v)| (*k, v.as_ref())).collect()
     }
 }
 
@@ -43,7 +40,7 @@ impl Group {
 
         for block in &self.blocks {
             let block_size = block.updates.len();
-            let one_large_batch = block.updates_flat_batch();
+            let one_large_batch = block.updates();
 
             group.throughput(Throughput::Elements(block_size as u64));
 
@@ -141,7 +138,7 @@ impl Benches {
                     Block {
                         smt: base_block
                             .smt
-                            .batch_update(base_block.updates_flat_batch(), &base_block.proof_reader)
+                            .batch_update(base_block.updates(), &base_block.proof_reader)
                             .unwrap(),
                         updates,
                         proof_reader,
@@ -167,8 +164,8 @@ impl Benches {
         rng: &mut StdRng,
         keys: &[HashValue],
         block_size: usize,
-    ) -> Vec<Vec<(HashValue, Option<StateValue>)>> {
-        std::iter::repeat_with(|| vec![Self::gen_update(rng, keys), Self::gen_update(rng, keys)])
+    ) -> Vec<(HashValue, Option<StateValue>)> {
+        std::iter::repeat_with(|| Self::gen_update(rng, keys))
             .take(block_size)
             .collect()
     }
@@ -188,11 +185,10 @@ impl Benches {
 
     fn gen_proof_reader(
         naive_smt: &mut NaiveSmt,
-        updates: &[Vec<(HashValue, Option<StateValue>)>],
+        updates: &[(HashValue, Option<StateValue>)],
     ) -> ProofReader {
         let proofs = updates
             .iter()
-            .flatten()
             .map(|(key, _)| (*key, naive_smt.get_proof(key)))
             .collect();
         ProofReader::new(proofs)