// This file is part of Substrate.
// Copyright (C) Parity Technologies (UK) Ltd.
// SPDX-License-Identifier: Apache-2.0
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use sc_cli::Result;
use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider};
use sc_client_db::{DbHash, DbState, DbStateBuilder};
use sp_blockchain::HeaderBackend;
use sp_database::{ColumnId, Transaction};
use sp_runtime::traits::{Block as BlockT, HashingFor, Header as HeaderT};
use sp_state_machine::Backend as StateBackend;
use sp_trie::PrefixedMemoryDB;
use log::{info, trace};
use rand::prelude::*;
use sp_storage::{ChildInfo, StateVersion};
use std::{
fmt::Debug,
sync::Arc,
time::{Duration, Instant},
};
use super::cmd::StorageCmd;
use crate::shared::{new_rng, BenchRecord};
impl StorageCmd {
/// Benchmarks the time it takes to write a single Storage item.
/// Uses the latest state that is available for the given client.
pub(crate) fn bench_write<Block, BA, H, C>(
&self,
client: Arc<C>,
(db, state_col): (Arc<dyn sp_database::Database<DbHash>>, ColumnId),
storage: Arc<dyn sp_state_machine::Storage<HashingFor<Block>>>,
shared_trie_cache: Option<sp_trie::cache::SharedTrieCache<HashingFor<Block>>>,
) -> Result<BenchRecord>
where
Block: BlockT<Header = H, Hash = DbHash> + Debug,
H: HeaderT<Hash = DbHash>,
BA: ClientBackend<Block>,
C: UsageProvider<Block> + HeaderBackend<Block> + StorageProvider<Block, BA>,
{
// Store the time that it took to write each value.
let mut record = BenchRecord::default();
let best_hash = client.usage_info().chain.best_hash;
let header = client.header(best_hash)?.ok_or("Header not found")?;
let original_root = *header.state_root();
info!("Preparing keys from block {}", best_hash);
let build_trie_backend =
|storage: Arc<dyn sp_state_machine::Storage<HashingFor<Block>>>,
original_root,
enable_pov_recorder: bool| {
let pov_recorder = enable_pov_recorder.then(|| Default::default());
DbStateBuilder::<HashingFor<Block>>::new(storage.clone(), original_root)
.with_optional_cache(shared_trie_cache.as_ref().map(|c| c.local_cache()))
.with_optional_recorder(pov_recorder)
.build()
};
let trie =
build_trie_backend(storage.clone(), original_root, !self.params.disable_pov_recorder);
// Load all KV pairs and randomly shuffle them.
let mut kvs: Vec<_> = trie.pairs(Default::default())?.collect();
let (mut rng, _) = new_rng(None);
kvs.shuffle(&mut rng);
info!("Writing {} keys", kvs.len());
let mut child_nodes = Vec::new();
let mut batched_keys = Vec::new();
// Generate all random values first; Make sure there are no collisions with existing
// db entries, so we can rollback all additions without corrupting existing entries.
for key_value in kvs {
let (k, original_v) = key_value?;
match (self.params.include_child_trees, self.is_child_key(k.to_vec())) {
(true, Some(info)) => {
let child_keys =
client.child_storage_keys(best_hash, info.clone(), None, None)?;
for ck in child_keys {
child_nodes.push((ck.clone(), info.clone()));
}
},
_ => {
// regular key
let mut new_v = vec![0; original_v.len()];
loop {
// Create a random value to overwrite with.
// NOTE: We use a possibly higher entropy than the original value,
// could be improved but acts as an over-estimation which is fine for now.
rng.fill_bytes(&mut new_v[..]);
if check_new_value::<Block>(
db.clone(),
&trie,
&k.to_vec(),
&new_v,
self.state_version(),
state_col,
None,
) {
break
}
}
batched_keys.push((k.to_vec(), new_v.to_vec()));
if batched_keys.len() < self.params.batch_size {
continue
}
// For every batched write use a different trie instance and recorder, so we
// don't benefit from past runs.
let trie = build_trie_backend(
storage.clone(),
original_root,
!self.params.disable_pov_recorder,
);
// Write each value in one commit.
let (size, duration) = measure_per_key_amortised_write_cost::<Block>(
db.clone(),
&trie,
batched_keys.clone(),
self.state_version(),
state_col,
None,
)?;
record.append(size, duration)?;
batched_keys.clear();
},
}
}
if self.params.include_child_trees {
child_nodes.shuffle(&mut rng);
info!("Writing {} child keys", child_nodes.len());
for (key, info) in child_nodes {
if let Some(original_v) = client
.child_storage(best_hash, &info.clone(), &key)
.expect("Checked above to exist")
{
let mut new_v = vec![0; original_v.0.len()];
loop {
rng.fill_bytes(&mut new_v[..]);
if check_new_value::<Block>(
db.clone(),
&trie,
&key.0,
&new_v,
self.state_version(),
state_col,
Some(&info),
) {
break
}
}
batched_keys.push((key.0, new_v.to_vec()));
if batched_keys.len() < self.params.batch_size {
continue
}
let trie = build_trie_backend(
storage.clone(),
original_root,
!self.params.disable_pov_recorder,
);
let (size, duration) = measure_per_key_amortised_write_cost::<Block>(
db.clone(),
&trie,
batched_keys.clone(),
self.state_version(),
state_col,
Some(&info),
)?;
record.append(size, duration)?;
batched_keys.clear();
}
}
}
Ok(record)
}
}
/// Converts a Trie transaction into a DB transaction.
/// Removals are ignored and will not be included in the final tx.
/// `invert_inserts` replaces all inserts with removals.
fn convert_tx<B: BlockT>(
db: Arc<dyn sp_database::Database<DbHash>>,
mut tx: PrefixedMemoryDB<HashingFor<B>>,
invert_inserts: bool,
col: ColumnId,
) -> Transaction<DbHash> {
let mut ret = Transaction::<DbHash>::default();
for (mut k, (v, rc)) in tx.drain().into_iter() {
if rc > 0 {
db.sanitize_key(&mut k);
if invert_inserts {
ret.remove(col, &k);
} else {
ret.set(col, &k, &v);
}
}
// < 0 means removal - ignored.
// 0 means no modification.
}
ret
}
/// Measures write benchmark
/// if `child_info` exist then it means this is a child tree key
fn measure_per_key_amortised_write_cost<Block: BlockT>(
db: Arc<dyn sp_database::Database<DbHash>>,
trie: &DbState<HashingFor<Block>>,
changes: Vec<(Vec<u8>, Vec<u8>)>,
version: StateVersion,
col: ColumnId,
child_info: Option<&ChildInfo>,
) -> Result<(usize, Duration)> {
let start = Instant::now();
// Create a TX that will modify the Trie in the DB and
// calculate the root hash of the Trie after the modification.
let average_len = changes.iter().map(|(_, v)| v.len()).sum::<usize>() / changes.len();
let replace = changes
.iter()
.map(|(key, new_v)| (key.as_ref(), Some(new_v.as_ref())))
.collect::<Vec<_>>();
let stx = match child_info {
Some(info) => trie.child_storage_root(info, replace.iter().cloned(), version).2,
None => trie.storage_root(replace.iter().cloned(), version).1,
};
// Only the keep the insertions, since we do not want to benchmark pruning.
let tx = convert_tx::<Block>(db.clone(), stx.clone(), false, col);
db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
let result = (average_len, start.elapsed() / changes.len() as u32);
// Now undo the changes by removing what was added.
let tx = convert_tx::<Block>(db.clone(), stx.clone(), true, col);
db.commit(tx).map_err(|e| format!("Writing to the Database: {}", e))?;
Ok(result)
}
/// Checks if a new value causes any collision in tree updates
/// returns true if there is no collision
/// if `child_info` exist then it means this is a child tree key
fn check_new_value<Block: BlockT>(
db: Arc<dyn sp_database::Database<DbHash>>,
trie: &DbState<HashingFor<Block>>,
key: &Vec<u8>,
new_v: &Vec<u8>,
version: StateVersion,
col: ColumnId,
child_info: Option<&ChildInfo>,
) -> bool {
let new_kv = vec![(key.as_ref(), Some(new_v.as_ref()))];
let mut stx = match child_info {
Some(info) => trie.child_storage_root(info, new_kv.iter().cloned(), version).2,
None => trie.storage_root(new_kv.iter().cloned(), version).1,
};
for (mut k, (_, rc)) in stx.drain().into_iter() {
if rc > 0 {
db.sanitize_key(&mut k);
if db.get(col, &k).is_some() {
trace!("Benchmark-store key creation: Key collision detected, retry");
return false
}
}
}
true
}