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
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, None, 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()
);
}