[diemdb_bench] initial_commit

Closes: aptos-labs#7407

Cargo.toml

@@ -129,6 +129,7 @@ members = [
     "storage/backup/backup-cli",
     "storage/backup/backup-service",
     "storage/diemdb",
+    "storage/diemdb-benchmark",
     "storage/diemsum",
     "storage/inspector",
     "storage/jellyfish-merkle",

storage/diemdb-benchmark/Cargo.toml

@@ -0,0 +1,27 @@
+[package]
+name = "diemdb-benchmark"
+version = "0.1.0"
+authors = ["Diem Association <[email protected]>"]
+description = "DiemDb benchmark"
+repository = "https://github.com/diem/diem"
+homepage = "https://diem.com"
+license = "Apache-2.0"
+publish = false
+edition = "2018"
+
+[dependencies]
+byteorder = "1.4.2"
+indicatif = "0.15.0"
+itertools = { version = "0.10.0", default-features = false }
+rand = "0.7.3"
+structopt = "0.3.21"
+
+diemdb = { path = "../diemdb", version = "0.1.0" }
+diem-config = { path = "../../config", version = "0.1.0" }
+diem-types = { path = "../../types", version = "0.1.0" }
+diem-workspace-hack = { path = "../../common/workspace-hack", version = "0.1.0" }
+storage-interface = { path = "../storage-interface", version = "0.1.0" }
+
+[features]
+default = []
+fuzzing = ["diem-types/fuzzing"]

storage/diemdb-benchmark/src/lib.rs

@@ -0,0 +1,102 @@
+// Copyright (c) The Diem Core Contributors
+// SPDX-License-Identifier: Apache-2.0
+
+use byteorder::{BigEndian, WriteBytesExt};
+use diem_config::config::RocksdbConfig;
+use diem_types::{
+    account_address::AccountAddress,
+    account_state_blob::AccountStateBlob,
+    transaction::{ChangeSet, Transaction, TransactionToCommit, WriteSetPayload},
+    vm_status::KeptVMStatus,
+    write_set::WriteSetMut,
+};
+use diemdb::DiemDB;
+use indicatif::{ProgressBar, ProgressStyle};
+use itertools::Itertools;
+use rand::Rng;
+use std::{collections::HashMap, fs, path::PathBuf};
+use storage_interface::DbWriter;
+
+fn gen_account_from_index(account_index: u64) -> AccountAddress {
+    let mut array = [0u8; AccountAddress::LENGTH];
+    array
+        .as_mut()
+        .write_u64::<BigEndian>(account_index)
+        .expect("Unable to write u64 to array");
+    AccountAddress::new(array)
+}
+
+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)
+}
+
+fn gen_txn_to_commit<R: Rng>(
+    max_accounts: u64,
+    blob_size: usize,
+    rng: &mut R,
+) -> TransactionToCommit {
+    let txn = Transaction::GenesisTransaction(WriteSetPayload::Direct(ChangeSet::new(
+        WriteSetMut::new(vec![])
+            .freeze()
+            .expect("freeze cannot fail"),
+        vec![],
+    )));
+    let account1 = gen_account_from_index(rng.gen_range(0, max_accounts));
+    let account2 = gen_account_from_index(rng.gen_range(0, max_accounts));
+    let mut states = HashMap::new();
+    let blob1 = gen_random_blob(blob_size, rng);
+    let blob2 = gen_random_blob(blob_size, rng);
+    states.insert(account1, blob1);
+    states.insert(account2, blob2);
+    TransactionToCommit::new(
+        txn,
+        states,
+        vec![], /* events */
+        0,      /* gas_used */
+        KeptVMStatus::Executed,
+    )
+}
+
+pub fn run_benchmark(num_accounts: usize, total_version: u64, blob_size: usize, db_dir: PathBuf) {
+    if db_dir.exists() {
+        fs::remove_dir_all(db_dir.join("diemdb")).unwrap();
+    }
+    // create if not exists
+    fs::create_dir_all(db_dir.clone()).unwrap();
+
+    let db = DiemDB::open(
+        &db_dir,
+        false, /* readonly */
+        None,  /* pruner */
+        RocksdbConfig::default(),
+    )
+    .expect("DB should open.");
+
+    let mut rng = ::rand::thread_rng();
+    let mut version = 0;
+
+    // Set a progressing bar
+    let bar = ProgressBar::new(total_version);
+    bar.set_style(
+        ProgressStyle::default_bar()
+            .template("[{elapsed}] {bar:100.cyan/blue} {pos:>7}/{len:7} {msg}"),
+    );
+
+    for chunk in &(0..total_version).chunks(1000 /* split by 1000 */) {
+        let txns_to_commit = chunk
+            .map(|_| gen_txn_to_commit(num_accounts as u64, blob_size, &mut rng))
+            .collect::<Vec<_>>();
+        let version_bump = txns_to_commit.len() as u64;
+        db.save_transactions(
+            &txns_to_commit,
+            version,
+            None, /* ledger_info_with_sigs */
+        )
+        .expect("commit cannot fail");
+        version = version.checked_add(version_bump).expect("Cannot overflow");
+        bar.inc(version_bump);
+    }
+    bar.finish();
+}

storage/diemdb-benchmark/src/main.rs

@@ -0,0 +1,25 @@
+// 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(short, default_value = "1000000")]
+    num_accounts: usize,
+
+    #[structopt(short, default_value = "1000000")]
+    version: 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();
+    diemdb_benchmark::run_benchmark(opt.num_accounts, opt.version, opt.blob_size, opt.db_dir);
+}

x.toml

@@ -125,6 +125,7 @@ members = [
     "diem-smoke-test-attribute",
     "diem-swarm",
     "diem-wallet",
+    "diemdb-benchmark",
     "executor-benchmark",
     "executor-test-helpers",
     "functional-tests",