use sc_cli::{CliConfiguration, DatabaseParams, PruningParams, Result, SharedParams};
use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider};
use sc_client_db::DbHash;
use sc_service::Configuration;
use sp_api::CallApiAt;
use sp_blockchain::HeaderBackend;
use sp_database::{ColumnId, Database};
use sp_runtime::traits::{Block as BlockT, HashingFor};
use sp_state_machine::Storage;
use sp_storage::{ChildInfo, ChildType, PrefixedStorageKey, StateVersion};
use clap::{Args, Parser, ValueEnum};
use log::info;
use serde::Serialize;
use sp_runtime::generic::BlockId;
use std::{fmt::Debug, path::PathBuf, sync::Arc};
use super::{
keys_selection::{select_entries, EmptyStorage as SelectEntriesEmptyStorage},
template::TemplateData,
};
use crate::shared::{HostInfoParams, WeightParams};
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Serialize, ValueEnum)]
pub enum StorageBenchmarkMode {
#[default]
ImportBlock,
ValidateBlock,
}
#[derive(Debug, Parser)]
pub struct StorageCmd {
#[allow(missing_docs)]
#[clap(flatten)]
pub shared_params: SharedParams,
#[allow(missing_docs)]
#[clap(flatten)]
pub database_params: DatabaseParams,
#[allow(missing_docs)]
#[clap(flatten)]
pub pruning_params: PruningParams,
#[allow(missing_docs)]
#[clap(flatten)]
pub params: StorageParams,
}
#[derive(Debug, Default, Serialize, Clone, PartialEq, Args)]
pub struct StorageParams {
#[allow(missing_docs)]
#[clap(flatten)]
pub weight_params: WeightParams,
#[allow(missing_docs)]
#[clap(flatten)]
pub hostinfo: HostInfoParams,
#[arg(long)]
pub skip_read: bool,
#[arg(long)]
pub skip_write: bool,
#[arg(long)]
pub template_path: Option<PathBuf>,
#[arg(long, value_name = "PATH")]
pub header: Option<PathBuf>,
#[arg(long)]
pub json_read_path: Option<PathBuf>,
#[arg(long)]
pub json_write_path: Option<PathBuf>,
#[arg(long, default_value_t = 1)]
pub warmups: u32,
#[arg(long, value_parser = clap::value_parser!(u8).range(0..=1))]
pub state_version: u8,
#[arg(long, value_name = "Bytes", default_value_t = 67108864)]
pub trie_cache_size: usize,
#[arg(long)]
pub enable_trie_cache: bool,
#[arg(long)]
pub include_child_trees: bool,
#[arg(long, default_value = "false")]
pub disable_pov_recorder: bool,
#[arg(long, default_value_t = 100_000)]
pub batch_size: usize,
#[arg(long, value_enum, default_value_t = StorageBenchmarkMode::ImportBlock)]
pub mode: StorageBenchmarkMode,
#[arg(long, default_value_t = 20)]
pub validate_block_rounds: u32,
#[arg(long)]
pub keys_limit: Option<usize>,
#[arg(long)]
pub child_keys_limit: Option<usize>,
#[arg(long)]
pub random_seed: Option<u64>,
}
impl StorageParams {
pub fn is_import_block_mode(&self) -> bool {
matches!(self.mode, StorageBenchmarkMode::ImportBlock)
}
pub fn is_validate_block_mode(&self) -> bool {
matches!(self.mode, StorageBenchmarkMode::ValidateBlock)
}
}
impl StorageCmd {
pub fn run<Block, BA, C>(
&self,
cfg: Configuration,
client: Arc<C>,
db: (Arc<dyn Database<DbHash>>, ColumnId),
storage: Arc<dyn Storage<HashingFor<Block>>>,
shared_trie_cache: Option<sp_trie::cache::SharedTrieCache<HashingFor<Block>>>,
) -> Result<()>
where
BA: ClientBackend<Block>,
Block: BlockT<Hash = DbHash>,
C: UsageProvider<Block>
+ StorageProvider<Block, BA>
+ HeaderBackend<Block>
+ CallApiAt<Block>,
{
let mut template = TemplateData::new(&cfg, &self.params)?;
let block_id = BlockId::<Block>::Number(client.usage_info().chain.best_number);
template.set_block_number(block_id.to_string());
if !self.params.skip_read {
self.bench_warmup(&client)?;
let record = self.bench_read(client.clone(), shared_trie_cache.clone())?;
if let Some(path) = &self.params.json_read_path {
record.save_json(&cfg, path, "read")?;
}
let stats = record.calculate_stats()?;
info!("Time summary [ns]:\n{:?}\nValue size summary:\n{:?}", stats.0, stats.1);
template.set_stats(Some(stats), None)?;
}
if !self.params.skip_write {
self.bench_warmup(&client)?;
let record = self.bench_write(client, db, storage, shared_trie_cache)?;
if let Some(path) = &self.params.json_write_path {
record.save_json(&cfg, path, "write")?;
}
let stats = record.calculate_stats()?;
info!("Time summary [ns]:\n{:?}\nValue size summary:\n{:?}", stats.0, stats.1);
template.set_stats(None, Some(stats))?;
}
template.write(&self.params.weight_params.weight_path, &self.params.template_path)
}
pub(crate) fn state_version(&self) -> StateVersion {
match self.params.state_version {
0 => StateVersion::V0,
1 => StateVersion::V1,
_ => unreachable!("Clap set to only allow 0 and 1"),
}
}
pub(crate) fn is_child_key(&self, key: Vec<u8>) -> Option<ChildInfo> {
if let Some((ChildType::ParentKeyId, storage_key)) =
ChildType::from_prefixed_key(&PrefixedStorageKey::new(key))
{
return Some(ChildInfo::new_default(storage_key));
}
None
}
fn bench_warmup<B, BA, C>(&self, client: &Arc<C>) -> Result<()>
where
C: UsageProvider<B> + StorageProvider<B, BA>,
B: BlockT + Debug,
BA: ClientBackend<B>,
{
let hash = client.usage_info().chain.best_hash;
let (keys, _) = select_entries(
self.params.keys_limit,
self.params.random_seed,
|first_key_ref| {
let fk = first_key_ref.map(|b| sp_storage::StorageKey(b.to_vec()));
Ok(client.storage_keys(hash, None, fk.as_ref())?)
},
|| Ok(client.storage_keys(hash, None, None)?),
|k: &sp_storage::StorageKey| k.0.as_slice(),
)?;
for i in 0..self.params.warmups {
info!("Warmup round {}/{}", i + 1, self.params.warmups);
let mut child_nodes = Vec::new();
for key in keys.as_slice() {
let _ = client
.storage(hash, &key)
.expect("Checked above to exist")
.ok_or("Value unexpectedly empty");
if let Some(info) = self
.params
.include_child_trees
.then(|| self.is_child_key(key.clone().0))
.flatten()
{
match select_entries(
self.params.child_keys_limit,
self.params.random_seed,
|first_key_ref| {
let fk = first_key_ref.map(|b| sp_storage::StorageKey(b.to_vec()));
Ok(client
.child_storage_keys(hash, info.clone(), None, fk.as_ref())?
.map(|ck| (ck, info.clone())))
},
|| {
Ok(client
.child_storage_keys(hash, info.clone(), None, None)?
.map(|ck| (ck, info.clone())))
},
|(k, _): &(sp_storage::StorageKey, sp_storage::ChildInfo)| k.0.as_slice(),
) {
Ok((entries, _)) => child_nodes.extend(entries),
Err(SelectEntriesEmptyStorage::Input(_)) => {},
Err(e) => return Err(e),
}
}
}
for (key, info) in child_nodes.as_slice() {
client
.child_storage(hash, info, key)
.expect("Checked above to exist")
.ok_or("Value unexpectedly empty")?;
}
}
Ok(())
}
}
impl CliConfiguration for StorageCmd {
fn shared_params(&self) -> &SharedParams {
&self.shared_params
}
fn database_params(&self) -> Option<&DatabaseParams> {
Some(&self.database_params)
}
fn pruning_params(&self) -> Option<&PruningParams> {
Some(&self.pruning_params)
}
fn trie_cache_maximum_size(&self) -> Result<Option<usize>> {
if self.params.enable_trie_cache && self.params.trie_cache_size > 0 {
Ok(Some(self.params.trie_cache_size))
} else {
Ok(None)
}
}
}