diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs b/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs index 307c9207fdaf9..f8a80db159a2c 100644 --- a/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs +++ b/substrate/utils/frame/benchmarking-cli/src/storage/cmd.rs @@ -116,6 +116,16 @@ pub struct StorageParams { /// Include child trees in benchmark. #[arg(long)] pub include_child_trees: bool, + + /// Maximum number of keys to read + /// (All keys if not define) + #[arg(long)] + pub keys_limit: Option, + + /// Seed to use for benchs randomness, the same seed allow to replay + /// becnhamrks under the same conditions. + #[arg(long)] + pub random_seed: Option, } impl StorageCmd { @@ -191,8 +201,14 @@ impl StorageCmd { BA: ClientBackend, { let hash = client.usage_info().chain.best_hash; - let mut keys: Vec<_> = client.storage_keys(hash, None, None)?.collect(); - let (mut rng, _) = new_rng(None); + let mut keys: Vec<_> = if let Some(keys_limit) = self.params.keys_limit { + use sp_core::blake2_256; + let first_key = self.params.random_seed.map(|seed| sp_storage::StorageKey(blake2_256(&seed.to_be_bytes()[..]).to_vec())); + client.storage_keys(hash, None, first_key.as_ref())?.take(keys_limit).collect() + } else { + client.storage_keys(hash, None, None)?.collect() + }; + let (mut rng, _) = new_rng(self.params.random_seed); keys.shuffle(&mut rng); for i in 0..self.params.warmups { diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/read.rs b/substrate/utils/frame/benchmarking-cli/src/storage/read.rs index fe72364269d50..d504d7f9e4ae7 100644 --- a/substrate/utils/frame/benchmarking-cli/src/storage/read.rs +++ b/substrate/utils/frame/benchmarking-cli/src/storage/read.rs @@ -40,9 +40,15 @@ impl StorageCmd { let best_hash = client.usage_info().chain.best_hash; info!("Preparing keys from block {}", best_hash); - // Load all keys and randomly shuffle them. - let mut keys: Vec<_> = client.storage_keys(best_hash, None, None)?.collect(); - let (mut rng, _) = new_rng(None); + // Load keys and randomly shuffle them. + let mut keys: Vec<_> = if let Some(keys_limit) = self.params.keys_limit { + use sp_core::blake2_256; + let first_key = self.params.random_seed.map(|seed| sp_storage::StorageKey(blake2_256(&seed.to_be_bytes()[..]).to_vec())); + client.storage_keys(best_hash, None, first_key.as_ref())?.take(keys_limit).collect() + } else { + client.storage_keys(best_hash, None, None)?.collect() + }; + let (mut rng, _) = new_rng(self.params.random_seed); keys.shuffle(&mut rng); let mut child_nodes = Vec::new(); diff --git a/substrate/utils/frame/benchmarking-cli/src/storage/write.rs b/substrate/utils/frame/benchmarking-cli/src/storage/write.rs index 4fa56b6e818f8..752883ed17090 100644 --- a/substrate/utils/frame/benchmarking-cli/src/storage/write.rs +++ b/substrate/utils/frame/benchmarking-cli/src/storage/write.rs @@ -60,9 +60,16 @@ impl StorageCmd { let trie = DbStateBuilder::>::new(storage.clone(), original_root).build(); info!("Preparing keys from block {}", best_hash); - // Load all KV pairs and randomly shuffle them. - let mut kvs: Vec<_> = trie.pairs(Default::default())?.collect(); - let (mut rng, _) = new_rng(None); + // Load KV pairs and randomly shuffle them. + let mut kvs: Vec<_> = if let Some(keys_limit) = self.params.keys_limit { + let start_at = self.params.random_seed.map(|seed| sp_core::blake2_256(&seed.to_be_bytes()[..]).to_vec()); + let mut iter_args = sp_state_machine::IterArgs::default(); + iter_args.start_at = start_at.as_deref(); + trie.pairs(iter_args)?.take(keys_limit).collect() + } else { + trie.pairs(Default::default())?.collect() + }; + let (mut rng, _) = new_rng(self.params.random_seed); kvs.shuffle(&mut rng); info!("Writing {} keys", kvs.len());