1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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 = diem_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()
    );
}