[scratchpad] benchmark tool

Closes: aptos-labs#8195

Cargo.toml

@@ -140,6 +140,7 @@ members = [
     "storage/jellyfish-merkle",
     "storage/schemadb",
     "storage/scratchpad",
+    "storage/scratchpad-benchmark",
     "storage/state-view",
     "storage/storage-client",
     "storage/storage-interface",

storage/diemdb-benchmark/Cargo.toml

@@ -18,6 +18,7 @@ rand = "0.8.3"
 structopt = "0.3.21"
 
 diemdb = { path = "../diemdb" }
+diem-crypto = { path = "../../crypto/crypto" }
 diem-jellyfish-merkle = { path = "../jellyfish-merkle" }
 diem-config = { path = "../../config" }
 diem-types = { path = "../../types" }

storage/diemdb-benchmark/src/lib.rs

@@ -3,13 +3,16 @@
 
 use byteorder::{BigEndian, WriteBytesExt};
 use diem_config::config::RocksdbConfig;
+use diem_crypto::hash::HashValue;
 use diem_jellyfish_merkle::metrics::{
     DIEM_JELLYFISH_INTERNAL_ENCODED_BYTES, DIEM_JELLYFISH_LEAF_ENCODED_BYTES,
     DIEM_JELLYFISH_STORAGE_READS,
 };
 use diem_types::{
     account_address::AccountAddress,
     account_state_blob::AccountStateBlob,
+    block_info::BlockInfo,
+    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
     transaction::{ChangeSet, Transaction, TransactionToCommit, WriteSetPayload},
     vm_status::KeptVMStatus,
     write_set::WriteSetMut,
@@ -20,10 +23,14 @@ use diemdb::{
 use indicatif::{ProgressBar, ProgressStyle};
 use itertools::Itertools;
 use rand::Rng;
-use std::{collections::HashMap, fs, path::PathBuf};
-use storage_interface::DbWriter;
+use std::{
+    collections::{BTreeMap, HashMap},
+    fs,
+    path::PathBuf,
+};
+use storage_interface::{DbReader, DbWriter};
 
-fn gen_account_from_index(account_index: u64) -> AccountAddress {
+pub fn gen_account_from_index(account_index: u64) -> AccountAddress {
     let mut array = [0u8; AccountAddress::LENGTH];
     array
         .as_mut()
@@ -32,7 +39,7 @@ fn gen_account_from_index(account_index: u64) -> AccountAddress {
     AccountAddress::new(array)
 }
 
-fn gen_random_blob<R: Rng>(size: usize, rng: &mut R) -> AccountStateBlob {
+pub fn gen_random_blob<R: Rng>(size: usize, rng: &mut R) -> AccountStateBlob {
     let mut v = vec![0u8; size];
     rng.fill(v.as_mut_slice());
     AccountStateBlob::from(v)
@@ -110,6 +117,23 @@ pub fn run_benchmark(
         version = version.checked_add(version_bump).expect("Cannot overflow");
         bar.inc(version_bump);
     }
+    let accu_root_hash = db.get_accumulator_root_hash(total_version - 1).unwrap();
+    // Last txn
+    let li = LedgerInfo::new(
+        BlockInfo::new(
+            /* current_epoch = */ 0,
+            /* round = */ 0,
+            /* block_id */ HashValue::random_with_rng(&mut rng),
+            accu_root_hash,
+            total_version - 1,
+            /* timestamp = */ 0,
+            None,
+        ),
+        HashValue::random_with_rng(&mut rng),
+    );
+    let li_with_sigs = LedgerInfoWithSignatures::new(li, BTreeMap::new());
+    db.save_transactions(&[], total_version, Some(&li_with_sigs))
+        .unwrap();
     bar.finish();
 
     db.update_rocksdb_properties().unwrap();

storage/scratchpad-benchmark/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "scratchpad-benchmark"
+version = "0.1.0"
+authors = ["Diem Association <[email protected]>"]
+description = "Diem executor benchmark"
+repository = "https://github.com/diem/diem"
+homepage = "https://diem.com"
+license = "Apache-2.0"
+publish = false
+edition = "2018"
+
+[dependencies]
+anyhow = "1.0.38"
+itertools = { version = "0.10.0", default-features = false }
+rand = "0.8.3"
+rayon = "1.5.0"
+structopt = "0.3.21"
+
+diemdb = { path = "../diemdb" }
+diemdb-benchmark = { path = "../diemdb-benchmark" }
+diem-config = { path = "../../config" }
+diem-crypto = { path = "../../crypto/crypto" }
+diem-genesis-tool = {path = "../../config/management/genesis", features = ["testing"] }
+diem-infallible = { path = "../../common/infallible" }
+diem-logger = { path = "../../common/logger" }
+diem-types = { path = "../../types" }
+diem-workspace-hack = { path = "../../common/workspace-hack" }
+executor-types = { path = "../../execution/executor-types" }
+scratchpad = { path = "..//scratchpad" }
+storage-interface = { path = "../storage-interface" }

storage/scratchpad-benchmark/src/lib.rs

@@ -0,0 +1,70 @@
+// Copyright (c) The Diem Core Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+use anyhow::Result;
+use diem_config::config::RocksdbConfig;
+use diem_types::{account_address::HashAccountAddress, account_state_blob::AccountStateBlob};
+use diemdb::DiemDB;
+use diemdb_benchmark::{gen_account_from_index, gen_random_blob};
+use executor_types::ProofReader;
+use rand::Rng;
+use std::{collections::HashMap, path::PathBuf};
+use storage_interface::DbReader;
+
+type SparseMerkleTree = scratchpad::SparseMerkleTree<AccountStateBlob>;
+
+pub fn run_benchmark(num_updates: usize, max_accounts: u64, blob_size: usize, db_dir: PathBuf) {
+    let db = DiemDB::open(
+        &db_dir,
+        false, /* readonly */
+        None,  /* pruner */
+        RocksdbConfig::default(),
+    )
+    .expect("DB should open.");
+
+    let mut rng = ::rand::thread_rng();
+
+    let updates = (0..num_updates)
+        .into_iter()
+        .map(|_| {
+            (
+                gen_account_from_index(rng.gen_range(0..max_accounts)),
+                gen_random_blob(blob_size, &mut rng),
+            )
+        })
+        .collect::<Vec<_>>();
+
+    let version = db.get_latest_version().unwrap();
+    let account_state_proofs = updates
+        .iter()
+        .map(|(k, _)| {
+            db.get_account_state_with_proof(*k, version, version)
+                .map(|p| p.proof.transaction_info_to_account_proof().clone())
+        })
+        .collect::<Result<Vec<_>>>()
+        .unwrap();
+
+    let proof_reader = ProofReader::new(
+        itertools::zip_eq(
+            updates.iter().map(|(k, _)| k.hash()),
+            account_state_proofs.into_iter(),
+        )
+        .collect::<HashMap<_, _>>(),
+    );
+    let root = db.get_latest_state_root().unwrap().1;
+    let smt = SparseMerkleTree::new(root);
+    let start = std::time::Instant::now();
+    smt.batch_update(
+        updates
+            .iter()
+            .map(|(k, v)| (k.hash(), v))
+            .collect::<Vec<_>>(),
+        &proof_reader,
+    )
+    .unwrap();
+    println!(
+        "Sparse Merkle Tree batch update {} updates: {}ms",
+        num_updates,
+        start.elapsed().as_millis()
+    );
+}

storage/scratchpad-benchmark/src/main.rs

@@ -0,0 +1,30 @@
+// Copyright (c) The Diem Core Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+use std::path::PathBuf;
+use structopt::StructOpt;
+
+#[derive(Debug, StructOpt)]
+struct Opt {
+    #[structopt(long, default_value = "2000")]
+    num_updates: usize,
+
+    #[structopt(long)]
+    num_accounts: u64,
+
+    #[structopt(short, default_value = "40")]
+    blob_size: usize,
+
+    #[structopt(long, parse(from_os_str))]
+    db_dir: PathBuf,
+}
+
+fn main() {
+    let opt = Opt::from_args();
+    scratchpad_benchmark::run_benchmark(
+        opt.num_updates,
+        opt.num_accounts,
+        opt.blob_size,
+        opt.db_dir,
+    );
+}

x.toml

@@ -190,6 +190,7 @@ members = [
     "move-prover-test-utils",
     "move-vm-integration-tests",
     "offchain",
+    "scratchpad-benchmark",
     "sdk-compatibility",
     "serializer-tests",
     "smoke-test",