From 523ddf5074fefb755af619cfa44825d39807a54e Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 31 Jan 2023 10:58:42 +0000 Subject: [PATCH 01/22] Rework storage iterators --- Cargo.lock | 40 +- Cargo.toml | 4 + client/api/src/backend.rs | 158 ++++--- client/db/src/bench.rs | 99 ++-- client/db/src/lib.rs | 82 ++-- client/db/src/record_stats_state.rs | 86 ++-- client/rpc/src/state/state_full.rs | 22 +- .../service/src/chain_ops/export_raw_state.rs | 46 +- client/service/src/client/client.rs | 62 +-- client/service/test/src/client/mod.rs | 25 +- primitives/state-machine/Cargo.toml | 6 +- primitives/state-machine/src/backend.rs | 236 +++++++++- primitives/state-machine/src/ext.rs | 55 ++- primitives/state-machine/src/lib.rs | 2 +- primitives/state-machine/src/testing.rs | 7 +- primitives/state-machine/src/trie_backend.rs | 83 ++-- .../state-machine/src/trie_backend_essence.rs | 426 ++++++++---------- primitives/storage/src/lib.rs | 1 + primitives/trie/Cargo.toml | 2 +- primitives/trie/src/lib.rs | 3 +- primitives/trie/src/recorder.rs | 1 + test-utils/runtime/Cargo.toml | 2 +- .../frame/benchmarking-cli/src/storage/cmd.rs | 4 +- .../benchmarking-cli/src/storage/read.rs | 7 +- .../benchmarking-cli/src/storage/write.rs | 10 +- .../rpc/state-trie-migration-rpc/Cargo.toml | 2 +- 26 files changed, 781 insertions(+), 690 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 29ddf38328519..2dafc0b5dbb26 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -139,9 +139,9 @@ dependencies = [ [[package]] name = "ahash" -version = "0.8.2" +version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf6ccdb167abbf410dcb915cabd428929d7f6a04980b54a11f26a39f1c7f7107" +checksum = "2c99f64d1e06488f620f932677e24bc6e2897582980441ae90a671415bd7ec2f" dependencies = [ "cfg-if", "getrandom 0.2.8", @@ -2847,8 +2847,7 @@ dependencies = [ [[package]] name = "hash-db" version = "0.15.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d23bd4e7b5eda0d0f3a307e8b381fdc8ba9000f26fbe912250c0a4cc3956364a" +source = "git+https://github.com/koute/trie.git?branch=master_iterators#a815264f668deb92a06c27233f5e899a60790049" [[package]] name = "hash256-std-hasher" @@ -2873,6 +2872,9 @@ name = "hashbrown" version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "43a3c133739dddd0d2990f9a4bdf8eb4b21ef50e4851ca85ab661199821d510e" +dependencies = [ + "ahash 0.8.3", +] [[package]] name = "heck" @@ -8385,7 +8387,7 @@ dependencies = [ name = "sc-finality-grandpa" version = "0.10.0-dev" dependencies = [ - "ahash 0.8.2", + "ahash 0.8.3", "array-bytes", "assert_matches", "async-trait", @@ -8585,7 +8587,7 @@ dependencies = [ name = "sc-network-gossip" version = "0.10.0-dev" dependencies = [ - "ahash 0.8.2", + "ahash 0.8.3", "futures", "futures-timer", "libp2p", @@ -9206,7 +9208,7 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "772575a524feeb803e5b0fcbc6dd9f367e579488197c94c6e4023aad2305774d" dependencies = [ - "ahash 0.8.2", + "ahash 0.8.3", "cfg-if", "hashbrown 0.13.2", ] @@ -10251,7 +10253,7 @@ dependencies = [ "sp-trie", "thiserror", "tracing", - "trie-db", + "trie-db 0.25.0", ] [[package]] @@ -10333,7 +10335,7 @@ dependencies = [ name = "sp-trie" version = "7.0.0" dependencies = [ - "ahash 0.8.2", + "ahash 0.8.3", "array-bytes", "criterion", "hash-db", @@ -10351,7 +10353,7 @@ dependencies = [ "thiserror", "tracing", "trie-bench", - "trie-db", + "trie-db 0.25.0", "trie-root", "trie-standardmap", ] @@ -10667,7 +10669,7 @@ dependencies = [ "sp-runtime", "sp-state-machine", "sp-trie", - "trie-db", + "trie-db 0.25.0", ] [[package]] @@ -10739,7 +10741,7 @@ dependencies = [ "sp-version", "substrate-test-runtime-client", "substrate-wasm-builder", - "trie-db", + "trie-db 0.25.0", ] [[package]] @@ -11322,7 +11324,7 @@ dependencies = [ "keccak-hasher", "memory-db", "parity-scale-codec", - "trie-db", + "trie-db 0.24.0", "trie-root", "trie-standardmap", ] @@ -11340,6 +11342,18 @@ dependencies = [ "smallvec", ] +[[package]] +name = "trie-db" +version = "0.25.0" +source = "git+https://github.com/koute/trie.git?branch=master_iterators#4d8bda6078eb911027ee3f43440a4272ab673ee5" +dependencies = [ + "hash-db", + "hashbrown 0.13.2", + "log", + "rustc-hex", + "smallvec", +] + [[package]] name = "trie-root" version = "0.17.0" diff --git a/Cargo.toml b/Cargo.toml index b004886e83f9a..a222a050c6107 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -321,3 +321,7 @@ inherits = "release" lto = "fat" # https://doc.rust-lang.org/rustc/codegen-options/index.html#codegen-units codegen-units = 1 + +[patch.crates-io] +trie-db = { git = "https://github.com/koute/trie.git", branch = "master_iterators" } +hash-db = { git = "https://github.com/koute/trie.git", branch = "master_iterators" } diff --git a/client/api/src/backend.rs b/client/api/src/backend.rs index 4ef609bdd4525..3880755f9908d 100644 --- a/client/api/src/backend.rs +++ b/client/api/src/backend.rs @@ -31,14 +31,13 @@ use sp_runtime::{ Justification, Justifications, StateVersion, Storage, }; use sp_state_machine::{ - backend::AsTrieBackend, ChildStorageCollection, IndexOperation, OffchainChangesCollection, - StorageCollection, + backend::AsTrieBackend, ChildStorageCollection, IndexOperation, IterArgs, + OffchainChangesCollection, StorageCollection, StorageIterator, }; use sp_storage::{ChildInfo, StorageData, StorageKey}; use std::collections::{HashMap, HashSet}; pub use sp_state_machine::{Backend as StateBackend, KeyValueStates}; -use std::marker::PhantomData; /// Extracts the state backend type for the given backend. pub type StateBackendFor = >::State; @@ -303,32 +302,50 @@ pub trait AuxStore { } /// An `Iterator` that iterates keys in a given block under a prefix. -pub struct KeyIterator { +pub struct KeysIter +where + State: StateBackend>, + Block: BlockT, +{ + inner: >>::RawIter, state: State, - child_storage: Option, - prefix: Option, - current_key: Vec, - _phantom: PhantomData, } -impl KeyIterator { - /// create a KeyIterator instance - pub fn new(state: State, prefix: Option, current_key: Vec) -> Self { - Self { state, child_storage: None, prefix, current_key, _phantom: PhantomData } +impl KeysIter +where + State: StateBackend>, + Block: BlockT, +{ + /// Create a new iterator over storage keys. + pub fn new( + state: State, + prefix: Option<&StorageKey>, + start_at: Option<&StorageKey>, + ) -> Result { + let mut args = IterArgs::default(); + args.prefix = prefix.as_ref().map(|prefix| prefix.0.as_slice()); + args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); + + Ok(Self { inner: state.raw_iter(args)?, state }) } - /// Create a `KeyIterator` instance for a child storage. + /// Create a new iterator over a child storage's keys. pub fn new_child( state: State, child_info: ChildInfo, - prefix: Option, - current_key: Vec, - ) -> Self { - Self { state, child_storage: Some(child_info), prefix, current_key, _phantom: PhantomData } + prefix: Option<&StorageKey>, + start_at: Option<&StorageKey>, + ) -> Result { + let mut args = IterArgs::default(); + args.prefix = prefix.as_ref().map(|prefix| prefix.0.as_slice()); + args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); + args.child_info = Some(child_info); + + Ok(Self { inner: state.raw_iter(args)?, state }) } } -impl Iterator for KeyIterator +impl Iterator for KeysIter where Block: BlockT, State: StateBackend>, @@ -336,25 +353,55 @@ where type Item = StorageKey; fn next(&mut self) -> Option { - let next_key = if let Some(child_info) = self.child_storage.as_ref() { - self.state.next_child_storage_key(child_info, &self.current_key) - } else { - self.state.next_storage_key(&self.current_key) - } - .ok() - .flatten()?; - // this terminates the iterator the first time it fails. - if let Some(ref prefix) = self.prefix { - if !next_key.starts_with(&prefix.0[..]) { - return None - } - } - self.current_key = next_key.clone(); - Some(StorageKey(next_key)) + self.inner.next_key(&self.state)?.ok().map(StorageKey) + } +} + +/// An `Iterator` that iterates keys and values in a given block under a prefix. +pub struct PairsIter +where + State: StateBackend>, + Block: BlockT, +{ + inner: >>::RawIter, + state: State, +} + +impl Iterator for PairsIter +where + Block: BlockT, + State: StateBackend>, +{ + type Item = (StorageKey, StorageData); + + fn next(&mut self) -> Option { + self.inner + .next_pair(&self.state)? + .ok() + .map(|(key, value)| (StorageKey(key), StorageData(value))) } } -/// Provides acess to storage primitives +impl PairsIter +where + State: StateBackend>, + Block: BlockT, +{ + /// Create a new iterator over storage key and value pairs. + pub fn new( + state: State, + prefix: Option<&StorageKey>, + start_at: Option<&StorageKey>, + ) -> Result { + let mut args = IterArgs::default(); + args.prefix = prefix.as_ref().map(|prefix| prefix.0.as_slice()); + args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); + + Ok(Self { inner: state.raw_iter(args)?, state }) + } +} + +/// Provides access to storage primitives pub trait StorageProvider> { /// Given a block's `Hash` and a key, return the value under the key in that block. fn storage( @@ -363,13 +410,6 @@ pub trait StorageProvider> { key: &StorageKey, ) -> sp_blockchain::Result>; - /// Given a block's `Hash` and a key prefix, return the matching storage keys in that block. - fn storage_keys( - &self, - hash: Block::Hash, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result>; - /// Given a block's `Hash` and a key, return the value under the hash in that block. fn storage_hash( &self, @@ -377,22 +417,23 @@ pub trait StorageProvider> { key: &StorageKey, ) -> sp_blockchain::Result>; - /// Given a block's `Hash` and a key prefix, return the matching child storage keys and values - /// in that block. - fn storage_pairs( + /// Given a block's `Hash` and a key prefix, returns a `KeysIter` iterates matching storage + /// keys in that block. + fn storage_keys( &self, hash: Block::Hash, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result>; + prefix: Option<&StorageKey>, + start_key: Option<&StorageKey>, + ) -> sp_blockchain::Result>; - /// Given a block's `Hash` and a key prefix, return a `KeyIterator` iterates matching storage - /// keys in that block. - fn storage_keys_iter( + /// Given a block's `Hash` and a key prefix, returns an iterator over the storage keys and + /// values in that block. + fn storage_pairs( &self, - hash: Block::Hash, + hash: ::Hash, prefix: Option<&StorageKey>, start_key: Option<&StorageKey>, - ) -> sp_blockchain::Result>; + ) -> sp_blockchain::Result>; /// Given a block's `Hash`, a key and a child storage key, return the value under the key in /// that block. @@ -403,24 +444,15 @@ pub trait StorageProvider> { key: &StorageKey, ) -> sp_blockchain::Result>; - /// Given a block's `Hash`, a key prefix, and a child storage key, return the matching child - /// storage keys. - fn child_storage_keys( - &self, - hash: Block::Hash, - child_info: &ChildInfo, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result>; - /// Given a block's `Hash` and a key `prefix` and a child storage key, - /// return a `KeyIterator` that iterates matching storage keys in that block. - fn child_storage_keys_iter( + /// returns a `KeysIter` that iterates matching storage keys in that block. + fn child_storage_keys( &self, hash: Block::Hash, child_info: ChildInfo, prefix: Option<&StorageKey>, start_key: Option<&StorageKey>, - ) -> sp_blockchain::Result>; + ) -> sp_blockchain::Result>; /// Given a block's `Hash`, a key and a child storage key, return the hash under the key in that /// block. diff --git a/client/db/src/bench.rs b/client/db/src/bench.rs index a744762876440..965999d0690dc 100644 --- a/client/db/src/bench.rs +++ b/client/db/src/bench.rs @@ -31,7 +31,8 @@ use sp_runtime::{ StateVersion, Storage, }; use sp_state_machine::{ - backend::Backend as StateBackend, ChildStorageCollection, DBValue, StorageCollection, + backend::Backend as StateBackend, ChildStorageCollection, DBValue, IterArgs, StorageCollection, + StorageIterator, StorageKey, StorageValue, }; use sp_trie::{ cache::{CacheSize, SharedTrieCache}, @@ -83,6 +84,31 @@ pub struct BenchmarkingState { shared_trie_cache: SharedTrieCache>, } +/// A raw iterator over the `BenchmarkingState`. +pub struct RawIter { + inner: as StateBackend>>::RawIter, +} + +impl StorageIterator> for RawIter { + type Backend = BenchmarkingState; + type Error = String; + + fn next_key(&mut self, backend: &Self::Backend) -> Option> { + self.inner.next_key(backend.state.borrow().as_ref()?) + } + + fn next_pair( + &mut self, + backend: &Self::Backend, + ) -> Option> { + self.inner.next_pair(backend.state.borrow().as_ref()?) + } + + fn was_complete(&self) -> bool { + self.inner.was_complete() + } +} + impl BenchmarkingState { /// Create a new instance that creates a database in a temporary dir. pub fn new( @@ -123,7 +149,7 @@ impl BenchmarkingState { ) }); let (root, transaction): (B::Hash, _) = - state.state.borrow_mut().as_mut().unwrap().full_storage_root( + state.state.borrow().as_ref().unwrap().full_storage_root( genesis.top.iter().map(|(k, v)| (k.as_ref(), Some(v.as_ref()))), child_delta, state_version, @@ -283,6 +309,7 @@ impl StateBackend> for BenchmarkingState { type Error = as StateBackend>>::Error; type Transaction = as StateBackend>>::Transaction; type TrieBackendStorage = as StateBackend>>::TrieBackendStorage; + type RawIter = RawIter; fn storage(&self, key: &[u8]) -> Result>, Self::Error> { self.add_read_key(None, key); @@ -356,58 +383,6 @@ impl StateBackend> for BenchmarkingState { .next_child_storage_key(child_info, key) } - fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { - if let Some(ref state) = *self.state.borrow() { - state.for_keys_with_prefix(prefix, f) - } - } - - fn for_key_values_with_prefix(&self, prefix: &[u8], f: F) { - if let Some(ref state) = *self.state.borrow() { - state.for_key_values_with_prefix(prefix, f) - } - } - - fn apply_to_key_values_while, Vec) -> bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - allow_missing: bool, - ) -> Result { - self.state.borrow().as_ref().ok_or_else(state_err)?.apply_to_key_values_while( - child_info, - prefix, - start_at, - f, - allow_missing, - ) - } - - fn apply_to_keys_while bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - ) { - if let Some(ref state) = *self.state.borrow() { - state.apply_to_keys_while(child_info, prefix, start_at, f) - } - } - - fn for_child_keys_with_prefix( - &self, - child_info: &ChildInfo, - prefix: &[u8], - f: F, - ) { - if let Some(ref state) = *self.state.borrow() { - state.for_child_keys_with_prefix(child_info, prefix, f) - } - } - fn storage_root<'a>( &self, delta: impl Iterator)>, @@ -437,19 +412,13 @@ impl StateBackend> for BenchmarkingState { .map_or(Default::default(), |s| s.child_storage_root(child_info, delta, state_version)) } - fn pairs(&self) -> Vec<(Vec, Vec)> { - self.state.borrow().as_ref().map_or(Default::default(), |s| s.pairs()) - } - - fn keys(&self, prefix: &[u8]) -> Vec> { - self.state.borrow().as_ref().map_or(Default::default(), |s| s.keys(prefix)) - } - - fn child_keys(&self, child_info: &ChildInfo, prefix: &[u8]) -> Vec> { + fn raw_iter(&self, args: IterArgs) -> Result { self.state .borrow() .as_ref() - .map_or(Default::default(), |s| s.child_keys(child_info, prefix)) + .map(|s| s.raw_iter(args)) + .unwrap_or(Ok(Default::default())) + .map(|raw_iter| RawIter { inner: raw_iter }) } fn commit( @@ -587,7 +556,7 @@ impl StateBackend> for BenchmarkingState { } fn register_overlay_stats(&self, stats: &sp_state_machine::StateMachineStats) { - self.state.borrow_mut().as_mut().map(|s| s.register_overlay_stats(stats)); + self.state.borrow().as_ref().map(|s| s.register_overlay_stats(stats)); } fn usage_info(&self) -> sp_state_machine::UsageInfo { diff --git a/client/db/src/lib.rs b/client/db/src/lib.rs index f217eb6480abc..b5e7d7a42342f 100644 --- a/client/db/src/lib.rs +++ b/client/db/src/lib.rs @@ -86,8 +86,9 @@ use sp_runtime::{ }; use sp_state_machine::{ backend::{AsTrieBackend, Backend as StateBackend}, - ChildStorageCollection, DBValue, IndexOperation, OffchainChangesCollection, StateMachineStats, - StorageCollection, UsageInfo as StateUsageInfo, + ChildStorageCollection, DBValue, IndexOperation, IterArgs, OffchainChangesCollection, + StateMachineStats, StorageCollection, StorageIterator, StorageKey, StorageValue, + UsageInfo as StateUsageInfo, }; use sp_trie::{cache::SharedTrieCache, prefixed_key, MemoryDB, PrefixedMemoryDB}; @@ -159,10 +160,36 @@ impl std::fmt::Debug for RefTrackingState { } } +/// A raw iterator over the `RefTrackingState`. +pub struct RawIter { + inner: as StateBackend>>::RawIter, +} + +impl StorageIterator> for RawIter { + type Backend = RefTrackingState; + type Error = as StateBackend>>::Error; + + fn next_key(&mut self, backend: &Self::Backend) -> Option> { + self.inner.next_key(&backend.state) + } + + fn next_pair( + &mut self, + backend: &Self::Backend, + ) -> Option> { + self.inner.next_pair(&backend.state) + } + + fn was_complete(&self) -> bool { + self.inner.was_complete() + } +} + impl StateBackend> for RefTrackingState { type Error = as StateBackend>>::Error; type Transaction = as StateBackend>>::Transaction; type TrieBackendStorage = as StateBackend>>::TrieBackendStorage; + type RawIter = RawIter; fn storage(&self, key: &[u8]) -> Result>, Self::Error> { self.state.storage(key) @@ -212,45 +239,6 @@ impl StateBackend> for RefTrackingState { self.state.next_child_storage_key(child_info, key) } - fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { - self.state.for_keys_with_prefix(prefix, f) - } - - fn for_key_values_with_prefix(&self, prefix: &[u8], f: F) { - self.state.for_key_values_with_prefix(prefix, f) - } - - fn apply_to_key_values_while, Vec) -> bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - allow_missing: bool, - ) -> Result { - self.state - .apply_to_key_values_while(child_info, prefix, start_at, f, allow_missing) - } - - fn apply_to_keys_while bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - ) { - self.state.apply_to_keys_while(child_info, prefix, start_at, f) - } - - fn for_child_keys_with_prefix( - &self, - child_info: &ChildInfo, - prefix: &[u8], - f: F, - ) { - self.state.for_child_keys_with_prefix(child_info, prefix, f) - } - fn storage_root<'a>( &self, delta: impl Iterator)>, @@ -274,16 +262,8 @@ impl StateBackend> for RefTrackingState { self.state.child_storage_root(child_info, delta, state_version) } - fn pairs(&self) -> Vec<(Vec, Vec)> { - self.state.pairs() - } - - fn keys(&self, prefix: &[u8]) -> Vec> { - self.state.keys(prefix) - } - - fn child_keys(&self, child_info: &ChildInfo, prefix: &[u8]) -> Vec> { - self.state.child_keys(child_info, prefix) + fn raw_iter(&self, args: IterArgs) -> Result { + self.state.raw_iter(args).map(|inner| RawIter { inner }) } fn register_overlay_stats(&self, stats: &StateMachineStats) { diff --git a/client/db/src/record_stats_state.rs b/client/db/src/record_stats_state.rs index 0b51d3fef2127..76d29db46d4d9 100644 --- a/client/db/src/record_stats_state.rs +++ b/client/db/src/record_stats_state.rs @@ -26,7 +26,7 @@ use sp_runtime::{ }; use sp_state_machine::{ backend::{AsTrieBackend, Backend as StateBackend}, - TrieBackend, + IterArgs, StorageIterator, StorageKey, StorageValue, TrieBackend, }; use std::sync::Arc; @@ -73,10 +73,43 @@ impl>, B: BlockT> RecordStatsState { } } +pub struct RawIter +where + S: StateBackend>, + B: BlockT, +{ + inner: >>::RawIter, +} + +impl StorageIterator> for RawIter +where + S: StateBackend>, + B: BlockT, +{ + type Backend = RecordStatsState; + type Error = S::Error; + + fn next_key(&mut self, backend: &Self::Backend) -> Option> { + self.inner.next_key(&backend.state) + } + + fn next_pair( + &mut self, + backend: &Self::Backend, + ) -> Option> { + self.inner.next_pair(&backend.state) + } + + fn was_complete(&self) -> bool { + self.inner.was_complete() + } +} + impl>, B: BlockT> StateBackend> for RecordStatsState { type Error = S::Error; type Transaction = S::Transaction; type TrieBackendStorage = S::TrieBackendStorage; + type RawIter = RawIter; fn storage(&self, key: &[u8]) -> Result>, Self::Error> { let value = self.state.storage(key)?; @@ -122,28 +155,6 @@ impl>, B: BlockT> StateBackend> for Record self.state.exists_child_storage(child_info, key) } - fn apply_to_key_values_while, Vec) -> bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - allow_missing: bool, - ) -> Result { - self.state - .apply_to_key_values_while(child_info, prefix, start_at, f, allow_missing) - } - - fn apply_to_keys_while bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - ) { - self.state.apply_to_keys_while(child_info, prefix, start_at, f) - } - fn next_storage_key(&self, key: &[u8]) -> Result>, Self::Error> { self.state.next_storage_key(key) } @@ -156,23 +167,6 @@ impl>, B: BlockT> StateBackend> for Record self.state.next_child_storage_key(child_info, key) } - fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { - self.state.for_keys_with_prefix(prefix, f) - } - - fn for_key_values_with_prefix(&self, prefix: &[u8], f: F) { - self.state.for_key_values_with_prefix(prefix, f) - } - - fn for_child_keys_with_prefix( - &self, - child_info: &ChildInfo, - prefix: &[u8], - f: F, - ) { - self.state.for_child_keys_with_prefix(child_info, prefix, f) - } - fn storage_root<'a>( &self, delta: impl Iterator)>, @@ -196,16 +190,8 @@ impl>, B: BlockT> StateBackend> for Record self.state.child_storage_root(child_info, delta, state_version) } - fn pairs(&self) -> Vec<(Vec, Vec)> { - self.state.pairs() - } - - fn keys(&self, prefix: &[u8]) -> Vec> { - self.state.keys(prefix) - } - - fn child_keys(&self, child_info: &ChildInfo, prefix: &[u8]) -> Vec> { - self.state.child_keys(child_info, prefix) + fn raw_iter(&self, args: IterArgs) -> Result { + self.state.raw_iter(args).map(|inner| RawIter { inner }) } fn register_overlay_stats(&self, stats: &sp_state_machine::StateMachineStats) { diff --git a/client/rpc/src/state/state_full.rs b/client/rpc/src/state/state_full.rs index 58dfd9ea876a1..d1aa959857c27 100644 --- a/client/rpc/src/state/state_full.rs +++ b/client/rpc/src/state/state_full.rs @@ -213,23 +213,29 @@ where .map_err(client_err) } + // TODO: This is horribly broken; either remove it, or make it streaming. fn storage_keys( &self, block: Option, prefix: StorageKey, ) -> std::result::Result, Error> { + // TODO: Remove the `.collect`. self.block_or_best(block) - .and_then(|block| self.client.storage_keys(block, &prefix)) + .and_then(|block| self.client.storage_keys(block, Some(&prefix), None)) + .map(|iter| iter.collect()) .map_err(client_err) } + // TODO: This is horribly broken; either remove it, or make it streaming. fn storage_pairs( &self, block: Option, prefix: StorageKey, ) -> std::result::Result, Error> { + // TODO: Remove the `.collect`. self.block_or_best(block) - .and_then(|block| self.client.storage_pairs(block, &prefix)) + .and_then(|block| self.client.storage_pairs(block, Some(&prefix), None)) + .map(|iter| iter.collect()) .map_err(client_err) } @@ -241,9 +247,7 @@ where start_key: Option, ) -> std::result::Result, Error> { self.block_or_best(block) - .and_then(|block| { - self.client.storage_keys_iter(block, prefix.as_ref(), start_key.as_ref()) - }) + .and_then(|block| self.client.storage_keys(block, prefix.as_ref(), start_key.as_ref())) .map(|iter| iter.take(count as usize).collect()) .map_err(client_err) } @@ -284,7 +288,7 @@ where } // The key doesn't point to anything, so it's probably a prefix. - let iter = match client.storage_keys_iter(block, Some(&key), None).map_err(client_err) { + let iter = match client.storage_keys(block, Some(&key), None).map_err(client_err) { Ok(iter) => iter, Err(e) => return Ok(Err(e)), }; @@ -536,6 +540,7 @@ where storage_key: PrefixedStorageKey, prefix: StorageKey, ) -> std::result::Result, Error> { + // TODO: Remove the `.collect`. self.block_or_best(block) .and_then(|block| { let child_info = match ChildType::from_prefixed_key(&storage_key) { @@ -543,8 +548,9 @@ where ChildInfo::new_default(storage_key), None => return Err(sp_blockchain::Error::InvalidChildStorageKey), }; - self.client.child_storage_keys(block, &child_info, &prefix) + self.client.child_storage_keys(block, child_info, Some(&prefix), None) }) + .map(|iter| iter.collect()) .map_err(client_err) } @@ -563,7 +569,7 @@ where ChildInfo::new_default(storage_key), None => return Err(sp_blockchain::Error::InvalidChildStorageKey), }; - self.client.child_storage_keys_iter( + self.client.child_storage_keys( block, child_info, prefix.as_ref(), diff --git a/client/service/src/chain_ops/export_raw_state.rs b/client/service/src/chain_ops/export_raw_state.rs index ca7a070086f45..a805c03fa4af6 100644 --- a/client/service/src/chain_ops/export_raw_state.rs +++ b/client/service/src/chain_ops/export_raw_state.rs @@ -21,7 +21,10 @@ use sc_client_api::{StorageProvider, UsageProvider}; use sp_core::storage::{well_known_keys, ChildInfo, Storage, StorageChild, StorageKey, StorageMap}; use sp_runtime::traits::Block as BlockT; -use std::{collections::HashMap, sync::Arc}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; /// Export the raw state at the given `block`. If `block` is `None`, the /// best block will be used. @@ -31,35 +34,30 @@ where B: BlockT, BA: sc_client_api::backend::Backend, { - let empty_key = StorageKey(Vec::new()); - let mut top_storage = client.storage_pairs(hash, &empty_key)?; + let mut top = BTreeMap::new(); let mut children_default = HashMap::new(); - // Remove all default child storage roots from the top storage and collect the child storage - // pairs. - while let Some(pos) = top_storage - .iter() - .position(|(k, _)| k.0.starts_with(well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX)) - { - let (key, _) = top_storage.swap_remove(pos); - - let key = - StorageKey(key.0[well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX.len()..].to_vec()); - let child_info = ChildInfo::new_default(&key.0); - - let keys = client.child_storage_keys(hash, &child_info, &empty_key)?; - let mut pairs = StorageMap::new(); - keys.into_iter().try_for_each(|k| { - if let Some(value) = client.child_storage(hash, &child_info, &k)? { - pairs.insert(k.0, value.0); + for (key, value) in client.storage_pairs(hash, None, None)? { + // Remove all default child storage roots from the top storage and collect the child storage + // pairs. + if key.0.starts_with(well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX) { + let child_root_key = StorageKey( + key.0[well_known_keys::DEFAULT_CHILD_STORAGE_KEY_PREFIX.len()..].to_vec(), + ); + let child_info = ChildInfo::new_default(&child_root_key.0); + let mut pairs = StorageMap::new(); + for child_key in client.child_storage_keys(hash, child_info.clone(), None, None)? { + if let Some(child_value) = client.child_storage(hash, &child_info, &child_key)? { + pairs.insert(child_key.0, child_value.0); + } } - Ok::<_, Error>(()) - })?; + children_default.insert(child_root_key.0, StorageChild { child_info, data: pairs }); + continue + } - children_default.insert(key.0, StorageChild { child_info, data: pairs }); + top.insert(key.0, value.0); } - let top = top_storage.into_iter().map(|(k, v)| (k.0, v.0)).collect(); Ok(Storage { top, children_default }) } diff --git a/client/service/src/client/client.rs b/client/service/src/client/client.rs index 6a75fad628681..01db050abfc9c 100644 --- a/client/service/src/client/client.rs +++ b/client/service/src/client/client.rs @@ -40,8 +40,8 @@ use sc_client_api::{ }, execution_extensions::ExecutionExtensions, notifications::{StorageEventStream, StorageNotifications}, - CallExecutor, ExecutorProvider, KeyIterator, OnFinalityAction, OnImportAction, ProofProvider, - UsageProvider, + CallExecutor, ExecutorProvider, KeysIter, OnFinalityAction, OnImportAction, PairsIter, + ProofProvider, UsageProvider, }; use sc_consensus::{ BlockCheckParams, BlockImportParams, ForkChoiceStrategy, ImportResult, StateAction, @@ -1465,52 +1465,37 @@ where Block: BlockT, { fn storage_keys( - &self, - hash: Block::Hash, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result> { - let keys = self.state_at(hash)?.keys(&key_prefix.0).into_iter().map(StorageKey).collect(); - Ok(keys) - } - - fn storage_pairs( &self, hash: ::Hash, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result> { + prefix: Option<&StorageKey>, + start_key: Option<&StorageKey>, + ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - let keys = state - .keys(&key_prefix.0) - .into_iter() - .map(|k| { - let d = state.storage(&k).ok().flatten().unwrap_or_default(); - (StorageKey(k), StorageData(d)) - }) - .collect(); - Ok(keys) + Ok(KeysIter::new(state, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) } - fn storage_keys_iter( + fn child_storage_keys( &self, hash: ::Hash, + child_info: ChildInfo, prefix: Option<&StorageKey>, start_key: Option<&StorageKey>, - ) -> sp_blockchain::Result> { + ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - let start_key = start_key.or(prefix).map(|key| key.0.clone()).unwrap_or_else(Vec::new); - Ok(KeyIterator::new(state, prefix.cloned(), start_key)) + Ok(KeysIter::new_child(state, child_info, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) } - fn child_storage_keys_iter( + fn storage_pairs( &self, hash: ::Hash, - child_info: ChildInfo, prefix: Option<&StorageKey>, start_key: Option<&StorageKey>, - ) -> sp_blockchain::Result> { + ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - let start_key = start_key.or(prefix).map(|key| key.0.clone()).unwrap_or_else(Vec::new); - Ok(KeyIterator::new_child(state, child_info, prefix.cloned(), start_key)) + Ok(PairsIter::new(state, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) } fn storage( @@ -1535,21 +1520,6 @@ where .map_err(|e| sp_blockchain::Error::from_state(Box::new(e))) } - fn child_storage_keys( - &self, - hash: ::Hash, - child_info: &ChildInfo, - key_prefix: &StorageKey, - ) -> sp_blockchain::Result> { - let keys = self - .state_at(hash)? - .child_keys(child_info, &key_prefix.0) - .into_iter() - .map(StorageKey) - .collect(); - Ok(keys) - } - fn child_storage( &self, hash: ::Hash, diff --git a/client/service/test/src/client/mod.rs b/client/service/test/src/client/mod.rs index 97c22a1cb509e..4b84d1496256e 100644 --- a/client/service/test/src/client/mod.rs +++ b/client/service/test/src/client/mod.rs @@ -1594,7 +1594,7 @@ fn returns_status_for_pruned_blocks() { } #[test] -fn storage_keys_iter_prefix_and_start_key_works() { +fn storage_keys_prefix_and_start_key_works() { let child_info = ChildInfo::new_default(b"child"); let client = TestClientBuilder::new() .add_extra_child_storage(&child_info, b"first".to_vec(), vec![0u8; 32]) @@ -1609,7 +1609,7 @@ fn storage_keys_iter_prefix_and_start_key_works() { let child_prefix = StorageKey(b"sec".to_vec()); let res: Vec<_> = client - .storage_keys_iter(block_hash, Some(&prefix), None) + .storage_keys(block_hash, Some(&prefix), None) .unwrap() .map(|x| x.0) .collect(); @@ -1623,7 +1623,7 @@ fn storage_keys_iter_prefix_and_start_key_works() { ); let res: Vec<_> = client - .storage_keys_iter( + .storage_keys( block_hash, Some(&prefix), Some(&StorageKey(array_bytes::hex2bytes_unchecked("3a636f6465"))), @@ -1634,7 +1634,7 @@ fn storage_keys_iter_prefix_and_start_key_works() { assert_eq!(res, [array_bytes::hex2bytes_unchecked("3a686561707061676573")]); let res: Vec<_> = client - .storage_keys_iter( + .storage_keys( block_hash, Some(&prefix), Some(&StorageKey(array_bytes::hex2bytes_unchecked("3a686561707061676573"))), @@ -1645,19 +1645,14 @@ fn storage_keys_iter_prefix_and_start_key_works() { assert_eq!(res, Vec::>::new()); let res: Vec<_> = client - .child_storage_keys_iter(block_hash, child_info.clone(), Some(&child_prefix), None) + .child_storage_keys(block_hash, child_info.clone(), Some(&child_prefix), None) .unwrap() .map(|x| x.0) .collect(); assert_eq!(res, [b"second".to_vec()]); let res: Vec<_> = client - .child_storage_keys_iter( - block_hash, - child_info, - None, - Some(&StorageKey(b"second".to_vec())), - ) + .child_storage_keys(block_hash, child_info, None, Some(&StorageKey(b"second".to_vec()))) .unwrap() .map(|x| x.0) .collect(); @@ -1665,7 +1660,7 @@ fn storage_keys_iter_prefix_and_start_key_works() { } #[test] -fn storage_keys_iter_works() { +fn storage_keys_works() { let client = substrate_test_runtime_client::new(); let block_hash = client.info().best_hash; @@ -1673,7 +1668,7 @@ fn storage_keys_iter_works() { let prefix = StorageKey(array_bytes::hex2bytes_unchecked("")); let res: Vec<_> = client - .storage_keys_iter(block_hash, Some(&prefix), None) + .storage_keys(block_hash, Some(&prefix), None) .unwrap() .take(9) .map(|x| array_bytes::bytes2hex("", &x.0)) @@ -1694,7 +1689,7 @@ fn storage_keys_iter_works() { ); let res: Vec<_> = client - .storage_keys_iter( + .storage_keys( block_hash, Some(&prefix), Some(&StorageKey(array_bytes::hex2bytes_unchecked("3a636f6465"))), @@ -1717,7 +1712,7 @@ fn storage_keys_iter_works() { ); let res: Vec<_> = client - .storage_keys_iter( + .storage_keys( block_hash, Some(&prefix), Some(&StorageKey(array_bytes::hex2bytes_unchecked( diff --git a/primitives/state-machine/Cargo.toml b/primitives/state-machine/Cargo.toml index 595dfe286d225..dd47c8e988478 100644 --- a/primitives/state-machine/Cargo.toml +++ b/primitives/state-machine/Cargo.toml @@ -16,7 +16,7 @@ targets = ["x86_64-unknown-linux-gnu"] [dependencies] codec = { package = "parity-scale-codec", version = "3.2.2", default-features = false } hash-db = { version = "0.15.2", default-features = false } -log = { version = "0.4.17", optional = true } +log = { version = "0.4.17", default-features = false } parking_lot = { version = "0.12.1", optional = true } rand = { version = "0.8.5", optional = true } smallvec = "1.8.0" @@ -33,7 +33,7 @@ array-bytes = "4.1" pretty_assertions = "1.2.1" rand = "0.8.5" sp-runtime = { version = "7.0.0", path = "../runtime" } -trie-db = "0.24.0" +trie-db = "0.25.0" assert_matches = "1.5" [features] @@ -41,7 +41,7 @@ default = ["std"] std = [ "codec/std", "hash-db/std", - "log", + "log/std", "parking_lot", "rand", "sp-core/std", diff --git a/primitives/state-machine/src/backend.rs b/primitives/state-machine/src/backend.rs index 791183c4d7e4d..551090fb1e228 100644 --- a/primitives/state-machine/src/backend.rs +++ b/primitives/state-machine/src/backend.rs @@ -24,12 +24,135 @@ use crate::{ StorageKey, StorageValue, UsageInfo, }; use codec::Encode; +use core::marker::PhantomData; use hash_db::Hasher; use sp_core::storage::{ChildInfo, StateVersion, TrackedStorageKey}; #[cfg(feature = "std")] use sp_core::traits::RuntimeCode; use sp_std::vec::Vec; +/// A struct containing arguments for iterating over the storage. +#[derive(Default)] +#[non_exhaustive] +pub struct IterArgs<'a> { + /// The prefix of the keys over which to iterate. + pub prefix: Option<&'a [u8]>, + + /// The key from which to start the iteration from. + pub start_at: Option<&'a [u8]>, + + /// The info of the child trie over which to iterate over. + pub child_info: Option, + + /// Whether to stop iteration when a missing trie node is reached. + /// + /// When a missing trie node is reached the iterator will: + /// - return an error if this is set to `false` (default) + /// - return `None` if this is set to `true` + pub allow_missing: bool, +} + +/// A trait for a raw storage iterator. +pub trait StorageIterator +where + H: Hasher, +{ + /// The state backend over which the iterator is iterating. + type Backend; + + /// The error type. + type Error; + + /// Fetches the next key from the storage. + fn next_key( + &mut self, + backend: &Self::Backend, + ) -> Option>; + + /// Fetches the next key and value from the storage. + fn next_pair( + &mut self, + backend: &Self::Backend, + ) -> Option>; + + /// Returns whether the trie end was reached. + /// + /// Will return `true` once we've reached the end of the trie. + fn was_complete(&self) -> bool; +} + +/// An iterator over storage keys and values. +pub struct PairsIter<'a, H, I> +where + H: Hasher, + I: StorageIterator, +{ + backend: Option<&'a I::Backend>, + raw_iter: I, + _phantom: PhantomData, +} + +impl<'a, H, I> Iterator for PairsIter<'a, H, I> +where + H: Hasher, + I: StorageIterator, +{ + type Item = Result<(Vec, Vec), >::Error>; + fn next(&mut self) -> Option { + self.raw_iter.next_pair(self.backend.as_ref()?) + } +} + +impl<'a, H, I> Default for PairsIter<'a, H, I> +where + H: Hasher, + I: StorageIterator + Default, +{ + fn default() -> Self { + Self { + backend: Default::default(), + raw_iter: Default::default(), + _phantom: Default::default(), + } + } +} + +/// An iterator over storage keys. +pub struct KeysIter<'a, H, I> +where + H: Hasher, + I: StorageIterator, +{ + backend: Option<&'a I::Backend>, + raw_iter: I, + _phantom: PhantomData, +} + +impl<'a, H, I> Iterator for KeysIter<'a, H, I> +where + H: Hasher, + I: StorageIterator, +{ + type Item = Result, >::Error>; + fn next(&mut self) -> Option { + self.raw_iter.next_key(self.backend.as_ref()?) + } +} + +impl<'a, H, I> Default for KeysIter<'a, H, I> +where + H: Hasher, + I: StorageIterator + Default, +{ + fn default() -> Self { + Self { + backend: Default::default(), + raw_iter: Default::default(), + _phantom: Default::default(), + } + } +} + /// A state backend is used to read state data and can have changes committed /// to it. /// @@ -44,6 +167,9 @@ pub trait Backend: sp_std::fmt::Debug { /// Type of trie backend storage. type TrieBackendStorage: TrieBackendStorage; + /// Type of the raw storage iterator. + type RawIter: StorageIterator; + /// Get keyed storage or None if there is nothing associated. fn storage(&self, key: &[u8]) -> Result, Self::Error>; @@ -95,43 +221,103 @@ pub trait Backend: sp_std::fmt::Debug { /// Otherwise an error is produced. /// /// Returns `true` if trie end is reached. + // TODO: Remove this. fn apply_to_key_values_while, Vec) -> bool>( &self, child_info: Option<&ChildInfo>, prefix: Option<&[u8]>, start_at: Option<&[u8]>, - f: F, + mut f: F, allow_missing: bool, - ) -> Result; + ) -> Result { + let args = IterArgs { + child_info: child_info.cloned(), + prefix, + start_at, + allow_missing, + ..IterArgs::default() + }; + let mut iter = self.pairs(args)?; + while let Some(key_value) = iter.next() { + let (key, value) = key_value?; + if !f(key, value) { + return Ok(false) + } + } + Ok(iter.raw_iter.was_complete()) + } /// Retrieve all entries keys of storage and call `f` for each of those keys. /// Aborts as soon as `f` returns false. + // TODO: Remove this. fn apply_to_keys_while bool>( &self, child_info: Option<&ChildInfo>, prefix: Option<&[u8]>, start_at: Option<&[u8]>, - f: F, - ); + mut f: F, + ) -> Result<(), Self::Error> { + let args = + IterArgs { child_info: child_info.cloned(), prefix, start_at, ..IterArgs::default() }; + + for key in self.keys(args)? { + if !f(&key?) { + return Ok(()) + } + } + Ok(()) + } /// Retrieve all entries keys which start with the given prefix and /// call `f` for each of those keys. - fn for_keys_with_prefix(&self, prefix: &[u8], mut f: F) { - self.for_key_values_with_prefix(prefix, |k, _v| f(k)) + // TODO: Remove this. + fn for_keys_with_prefix( + &self, + prefix: &[u8], + mut f: F, + ) -> Result<(), Self::Error> { + let args = IterArgs { prefix: Some(prefix), ..IterArgs::default() }; + self.keys(args)?.try_for_each(|key| { + f(&key?); + Ok(()) + }) } /// Retrieve all entries keys and values of which start with the given prefix and /// call `f` for each of those keys. - fn for_key_values_with_prefix(&self, prefix: &[u8], f: F); + // TODO: Remove this. + fn for_key_values_with_prefix( + &self, + prefix: &[u8], + mut f: F, + ) -> Result<(), Self::Error> { + let args = IterArgs { prefix: Some(prefix), ..IterArgs::default() }; + self.pairs(args)?.try_for_each(|key_value| { + let (key, value) = key_value?; + f(&key, &value); + Ok(()) + }) + } /// Retrieve all child entries keys which start with the given prefix and /// call `f` for each of those keys. + // TODO: Remove this. fn for_child_keys_with_prefix( &self, child_info: &ChildInfo, prefix: &[u8], - f: F, - ); + mut f: F, + ) -> Result<(), Self::Error> { + let args = IterArgs { + child_info: Some(child_info.clone()), + prefix: Some(prefix), + ..IterArgs::default() + }; + self.keys(args)?.try_for_each(|key| { + f(&key?); + Ok(()) + }) + } /// Calculate the storage root, with given delta over what is already stored in /// the backend, and produce a "transaction" that can be used to commit. @@ -156,21 +342,25 @@ pub trait Backend: sp_std::fmt::Debug { where H::Out: Ord; - /// Get all key/value pairs into a Vec. - fn pairs(&self) -> Vec<(StorageKey, StorageValue)>; + /// Returns a lifetimeless raw storage iterator. + fn raw_iter(&self, args: IterArgs) -> Result; - /// Get all keys with given prefix - fn keys(&self, prefix: &[u8]) -> Vec { - let mut all = Vec::new(); - self.for_keys_with_prefix(prefix, |k| all.push(k.to_vec())); - all + /// Get an iterator over key/value pairs. + fn pairs<'a>(&'a self, args: IterArgs) -> Result, Self::Error> { + Ok(PairsIter { + backend: Some(self), + raw_iter: self.raw_iter(args)?, + _phantom: Default::default(), + }) } - /// Get all keys of child storage with given prefix - fn child_keys(&self, child_info: &ChildInfo, prefix: &[u8]) -> Vec { - let mut all = Vec::new(); - self.for_child_keys_with_prefix(child_info, prefix, |k| all.push(k.to_vec())); - all + /// Get an iterator over keys. + fn keys<'a>(&'a self, args: IterArgs) -> Result, Self::Error> { + Ok(KeysIter { + backend: Some(self), + raw_iter: self.raw_iter(args)?, + _phantom: Default::default(), + }) } /// Calculate the storage root, with given delta over what is already stored @@ -309,7 +499,7 @@ where #[cfg(feature = "std")] pub struct BackendRuntimeCode<'a, B, H> { backend: &'a B, - _marker: std::marker::PhantomData, + _marker: PhantomData, } #[cfg(feature = "std")] @@ -332,7 +522,7 @@ where { /// Create a new instance. pub fn new(backend: &'a B) -> Self { - Self { backend, _marker: std::marker::PhantomData } + Self { backend, _marker: PhantomData } } /// Return the [`RuntimeCode`] build from the wrapped `backend`. diff --git a/primitives/state-machine/src/ext.rs b/primitives/state-machine/src/ext.rs index 1db0ec517015b..43059f334336e 100644 --- a/primitives/state-machine/src/ext.rs +++ b/primitives/state-machine/src/ext.rs @@ -159,9 +159,10 @@ where use std::collections::HashMap; self.backend - .pairs() - .iter() - .map(|&(ref k, ref v)| (k.to_vec(), Some(v.to_vec()))) + .pairs(Default::default()) + .expect("never fails in tests; qed.") + .map(|key_value| key_value.expect("never fails in tests; qed.")) + .map(|(k, v)| (k, Some(v))) .chain(self.overlay.changes().map(|(k, v)| (k.clone(), v.value().cloned()))) .collect::>() .into_iter() @@ -757,28 +758,34 @@ where let mut delete_count: u32 = 0; let mut loop_count: u32 = 0; let mut maybe_next_key = None; - self.backend - .apply_to_keys_while(maybe_child, maybe_prefix, maybe_cursor, |key| { - if maybe_limit.map_or(false, |limit| loop_count == limit) { - maybe_next_key = Some(key.to_vec()); - return false - } - let overlay = match maybe_child { - Some(child_info) => self.overlay.child_storage(child_info, key), - None => self.overlay.storage(key), - }; - if !matches!(overlay, Some(None)) { - // not pending deletion from the backend - delete it. - if let Some(child_info) = maybe_child { - self.overlay.set_child_storage(child_info, key.to_vec(), None); - } else { - self.overlay.set_storage(key.to_vec(), None); + let result = + self.backend + .apply_to_keys_while(maybe_child, maybe_prefix, maybe_cursor, |key| { + if maybe_limit.map_or(false, |limit| loop_count == limit) { + maybe_next_key = Some(key.to_vec()); + return false } - delete_count = delete_count.saturating_add(1); - } - loop_count = loop_count.saturating_add(1); - true - }); + let overlay = match maybe_child { + Some(child_info) => self.overlay.child_storage(child_info, key), + None => self.overlay.storage(key), + }; + if !matches!(overlay, Some(None)) { + // not pending deletion from the backend - delete it. + if let Some(child_info) = maybe_child { + self.overlay.set_child_storage(child_info, key.to_vec(), None); + } else { + self.overlay.set_storage(key.to_vec(), None); + } + delete_count = delete_count.saturating_add(1); + } + loop_count = loop_count.saturating_add(1); + true + }); + + if let Err(error) = result { + log::debug!(target: "trie", "Error while iterating the storage: {}", error); + } + (maybe_next_key, delete_count, loop_count) } } diff --git a/primitives/state-machine/src/lib.rs b/primitives/state-machine/src/lib.rs index 188fb2f4b8d99..44a0daf6d4c6e 100644 --- a/primitives/state-machine/src/lib.rs +++ b/primitives/state-machine/src/lib.rs @@ -123,7 +123,7 @@ impl sp_std::fmt::Display for DefaultError { } pub use crate::{ - backend::Backend, + backend::{Backend, IterArgs, KeysIter, PairsIter, StorageIterator}, error::{Error, ExecutionError}, ext::Ext, overlayed_changes::{ diff --git a/primitives/state-machine/src/testing.rs b/primitives/state-machine/src/testing.rs index e94d34b5560cd..0e5e58bbecaae 100644 --- a/primitives/state-machine/src/testing.rs +++ b/primitives/state-machine/src/testing.rs @@ -241,7 +241,12 @@ where H::Out: Ord + codec::Codec, { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "overlay: {:?}\nbackend: {:?}", self.overlay, self.backend.pairs()) + let pairs: Vec<_> = self + .backend + .pairs(Default::default()) + .expect("creating an iterator over all of the pairs doesn't fail in tests") + .collect(); // TODO: Print those out without a `collect`. + write!(f, "overlay: {:?}\nbackend: {:?}", self.overlay, pairs) } } diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index ff5cce24c5124..312933716b675 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -20,6 +20,7 @@ #[cfg(feature = "std")] use crate::backend::AsTrieBackend; use crate::{ + backend::IterArgs, trie_backend_essence::{TrieBackendEssence, TrieBackendStorage}, Backend, StorageKey, StorageValue, }; @@ -28,7 +29,6 @@ use codec::Codec; use hash_db::HashDB; use hash_db::Hasher; use sp_core::storage::{ChildInfo, StateVersion}; -use sp_std::vec::Vec; #[cfg(feature = "std")] use sp_trie::{cache::LocalTrieCache, recorder::Recorder}; #[cfg(feature = "std")] @@ -51,6 +51,7 @@ pub trait AsLocalTrieCache: sealed::Sealed { impl AsLocalTrieCache for LocalTrieCache { #[cfg(feature = "std")] + #[inline] fn as_local_trie_cache(&self) -> &LocalTrieCache { self } @@ -58,6 +59,7 @@ impl AsLocalTrieCache for LocalTrieCache { #[cfg(feature = "std")] impl AsLocalTrieCache for &LocalTrieCache { + #[inline] fn as_local_trie_cache(&self) -> &LocalTrieCache { self } @@ -236,6 +238,7 @@ where type Error = crate::DefaultError; type Transaction = S::Overlay; type TrieBackendStorage = S; + type RawIter = crate::trie_backend_essence::RawIter; fn storage_hash(&self, key: &[u8]) -> Result, Self::Error> { self.essence.storage_hash(key) @@ -273,51 +276,8 @@ where self.essence.next_child_storage_key(child_info, key) } - fn for_keys_with_prefix(&self, prefix: &[u8], f: F) { - self.essence.for_keys_with_prefix(prefix, f) - } - - fn for_key_values_with_prefix(&self, prefix: &[u8], f: F) { - self.essence.for_key_values_with_prefix(prefix, f) - } - - fn apply_to_key_values_while, Vec) -> bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - allow_missing: bool, - ) -> Result { - self.essence - .apply_to_key_values_while(child_info, prefix, start_at, f, allow_missing) - } - - fn apply_to_keys_while bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - ) { - self.essence.apply_to_keys_while(child_info, prefix, start_at, f) - } - - fn for_child_keys_with_prefix( - &self, - child_info: &ChildInfo, - prefix: &[u8], - f: F, - ) { - self.essence.for_child_keys_with_prefix(child_info, prefix, f) - } - - fn pairs(&self) -> Vec<(StorageKey, StorageValue)> { - self.essence.pairs() - } - - fn keys(&self, prefix: &[u8]) -> Vec { - self.essence.keys(prefix) + fn raw_iter(&self, args: IterArgs) -> Result { + self.essence.raw_iter(args) } fn storage_root<'a>( @@ -579,7 +539,11 @@ pub mod tests { cache: Option, recorder: Option, ) { - assert!(!test_trie(state_version, cache, recorder).pairs().is_empty()); + assert!(!test_trie(state_version, cache, recorder) + .pairs(Default::default()) + .unwrap() + .next() + .is_none()); } #[test] @@ -589,8 +553,10 @@ pub mod tests { Default::default(), ) .build() - .pairs() - .is_empty()); + .pairs(Default::default()) + .unwrap() + .next() + .is_none()); } parameterized_test!(storage_root_is_non_default, storage_root_is_non_default_inner); @@ -638,7 +604,8 @@ pub mod tests { trie.for_keys_with_prefix(b"value", |key| { let for_first_time = seen.insert(key.to_vec()); assert!(for_first_time, "Seen key '{:?}' more than once", key); - }); + }) + .unwrap(); let mut expected = HashSet::new(); expected.insert(b"value1".to_vec()); @@ -664,7 +631,8 @@ pub mod tests { .collect::>(); let trie = test_trie(state_version, cache, recorder); - let keys = trie.keys(&[]); + let keys: Vec<_> = + trie.keys(Default::default()).unwrap().map(|result| result.unwrap()).collect(); assert_eq!(expected, keys); } @@ -724,7 +692,18 @@ pub mod tests { .with_recorder(Recorder::default()) .build(); assert_eq!(trie_backend.storage(b"key").unwrap(), proving_backend.storage(b"key").unwrap()); - assert_eq!(trie_backend.pairs(), proving_backend.pairs()); + assert_eq!( + trie_backend + .pairs(Default::default()) + .unwrap() + .map(|result| result.unwrap()) + .collect::>(), + proving_backend + .pairs(Default::default()) + .unwrap() + .map(|result| result.unwrap()) + .collect::>() + ); let (trie_root, mut trie_mdb) = trie_backend.storage_root(std::iter::empty(), state_version); diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index cdd1bb0bba055..7c12c86127e96 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -19,24 +19,23 @@ //! from storage. use crate::{ - backend::Consolidate, debug, trie_backend::AsLocalTrieCache, warn, StorageKey, StorageValue, + backend::{Consolidate, IterArgs, StorageIterator}, + trie_backend::AsLocalTrieCache, + warn, StorageKey, StorageValue, }; use codec::Codec; use hash_db::{self, AsHashDB, HashDB, HashDBRef, Hasher, Prefix}; #[cfg(feature = "std")] use parking_lot::RwLock; use sp_core::storage::{ChildInfo, ChildType, StateVersion}; -#[cfg(not(feature = "std"))] -use sp_std::marker::PhantomData; -use sp_std::{boxed::Box, vec::Vec}; +use sp_std::{boxed::Box, marker::PhantomData, vec::Vec}; #[cfg(feature = "std")] use sp_trie::recorder::Recorder; use sp_trie::{ child_delta_trie_root, delta_trie_root, empty_child_trie_root, read_child_trie_hash, read_child_trie_value, read_trie_value, trie_types::{TrieDBBuilder, TrieError}, - DBValue, KeySpacedDB, NodeCodec, Trie, TrieCache, TrieDBIterator, TrieDBKeyIterator, - TrieRecorder, + DBValue, KeySpacedDB, NodeCodec, Trie, TrieCache, TrieDBRawIterator, TrieRecorder, }; #[cfg(feature = "std")] use std::{collections::HashMap, sync::Arc}; @@ -257,6 +256,166 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss } } +enum IterState { + Pending, + FinishedComplete, + FinishedIncomplete, +} + +pub struct RawIter +where + H: Hasher, +{ + allow_missing: bool, + root: H::Out, + child_info: Option, + trie_iter: TrieDBRawIterator>, + state: IterState, + _phantom: PhantomData<(S, C)>, +} + +impl RawIter +where + H: Hasher, + S: TrieBackendStorage, + H::Out: Codec + Ord, + C: AsLocalTrieCache + Send + Sync, +{ + #[inline] + fn prepare( + &mut self, + backend: &TrieBackendEssence, + callback: impl FnOnce( + &sp_trie::TrieDB>, + &mut TrieDBRawIterator>, + ) -> Option::Out>>>>, + ) -> Option> { + if !matches!(self.state, IterState::Pending) { + return None + } + + let result = backend.with_trie_db(self.root, self.child_info.as_ref(), |db| { + callback(&db, &mut self.trie_iter) + }); + match result { + Some(Ok(key_value)) => Some(Ok(key_value)), + None => { + self.state = IterState::FinishedComplete; + None + }, + Some(Err(error)) => { + self.state = IterState::FinishedIncomplete; + if matches!(*error, TrieError::IncompleteDatabase(_)) && self.allow_missing { + None + } else { + Some(Err(format!("TrieDB iteration error: {}", error))) + } + }, + } + } +} + +impl Default for RawIter +where + H: Hasher, +{ + fn default() -> Self { + Self { + allow_missing: false, + child_info: None, + root: Default::default(), + trie_iter: TrieDBRawIterator::empty(), + state: IterState::FinishedComplete, + _phantom: Default::default(), + } + } +} + +impl TrieBackendEssence +where + H: Hasher, + H::Out: Codec + Ord, + S: TrieBackendStorage, + C: AsLocalTrieCache + Send + Sync, +{ + #[cfg(not(feature = "std"))] + #[inline] + fn with_trie_db( + &self, + root: H::Out, + child_info: Option<&ChildInfo>, + callback: impl FnOnce(&sp_trie::TrieDB>) -> R, + ) -> R { + let backend = self as &dyn HashDBRef>; + let db = child_info + .as_ref() + .map(|child_info| KeySpacedDB::new(backend, child_info.keyspace())); + let db = db.as_ref().map(|db| db as &dyn HashDBRef>).unwrap_or(backend); + + let trie = TrieDBBuilder::::new(db, &root).build(); + + callback(&trie) + } + + #[cfg(feature = "std")] + #[inline] + fn with_trie_db( + &self, + root: H::Out, + child_info: Option<&ChildInfo>, + callback: impl FnOnce(&sp_trie::TrieDB>) -> R, + ) -> R { + let backend = self as &dyn HashDBRef>; + let db = child_info + .as_ref() + .map(|child_info| KeySpacedDB::new(backend, child_info.keyspace())); + let db = db.as_ref().map(|db| db as &dyn HashDBRef>).unwrap_or(backend); + + let recorder = self.recorder.as_ref(); + let mut recorder = recorder.as_ref().map(|r| r.as_trie_recorder(root)); + let recorder = match recorder.as_mut() { + Some(recorder) => Some(recorder as &mut dyn TrieRecorder), + None => None, + }; + + let cache = self.trie_node_cache.as_ref(); + let mut cache = cache.map(|c| c.as_local_trie_cache().as_trie_db_cache(root)); + let cache = cache.as_mut().map(|c| c as &mut dyn sp_trie::TrieCache>); + + let trie = TrieDBBuilder::::new(db, &root) + .with_optional_recorder(recorder) + .with_optional_cache(cache) + .build(); + + callback(&trie) + } +} + +impl StorageIterator for RawIter +where + H: Hasher, + S: TrieBackendStorage, + H::Out: Codec + Ord, + C: AsLocalTrieCache + Send + Sync, +{ + type Backend = crate::TrieBackend; + type Error = crate::DefaultError; + + #[inline] + fn next_key(&mut self, backend: &Self::Backend) -> Option> { + self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_key(&trie)) + } + + #[inline] + fn next_pair(&mut self, backend: &Self::Backend) -> Option> { + self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_item(&trie)) + } + + fn was_complete(&self) -> bool { + matches!(self.state, IterState::FinishedComplete) + } +} + impl, H: Hasher, C: AsLocalTrieCache + Send + Sync> TrieBackendEssence where @@ -429,246 +588,41 @@ where }) } - /// Retrieve all entries keys of storage and call `f` for each of those keys. - /// Aborts as soon as `f` returns false. - /// - /// Returns `true` when all keys were iterated. - pub fn apply_to_key_values_while( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: impl FnMut(Vec, Vec) -> bool, - allow_missing_nodes: bool, - ) -> Result { - let root = if let Some(child_info) = child_info.as_ref() { - match self.child_root(child_info)? { - Some(child_root) => child_root, - None => return Ok(true), - } - } else { - self.root - }; - - self.trie_iter_inner(&root, prefix, f, child_info, start_at, allow_missing_nodes) - } - - /// Retrieve all entries keys of a storage and call `f` for each of those keys. - /// Aborts as soon as `f` returns false. - pub fn apply_to_keys_while bool>( - &self, - child_info: Option<&ChildInfo>, - prefix: Option<&[u8]>, - start_at: Option<&[u8]>, - f: F, - ) { - let root = if let Some(child_info) = child_info.as_ref() { - match self.child_root(child_info) { - Ok(Some(v)) => v, - // If the child trie doesn't exist, there is no need to continue. - Ok(None) => return, - Err(e) => { - debug!(target: "trie", "Error while iterating child storage: {}", e); - return - }, - } + pub fn raw_iter(&self, args: IterArgs) -> Result> { + let root = if let Some(child_info) = args.child_info.as_ref() { + let root = match self.child_root(&child_info)? { + Some(root) => root, + None => return Ok(Default::default()), + }; + root } else { self.root }; - self.trie_iter_key_inner(&root, prefix, f, child_info, start_at) - } - - /// Execute given closure for all keys starting with prefix. - pub fn for_child_keys_with_prefix( - &self, - child_info: &ChildInfo, - prefix: &[u8], - mut f: impl FnMut(&[u8]), - ) { - let root = match self.child_root(child_info) { - Ok(Some(v)) => v, - // If the child trie doesn't exist, there is no need to continue. - Ok(None) => return, - Err(e) => { - debug!(target: "trie", "Error while iterating child storage: {}", e); - return - }, - }; - - self.trie_iter_key_inner( - &root, - Some(prefix), - |k| { - f(k); - true - }, - Some(child_info), - None, - ) - } - - /// Execute given closure for all keys starting with prefix. - pub fn for_keys_with_prefix(&self, prefix: &[u8], mut f: F) { - self.trie_iter_key_inner( - &self.root, - Some(prefix), - |k| { - f(k); - true - }, - None, - None, - ) - } - - fn trie_iter_key_inner bool>( - &self, - root: &H::Out, - maybe_prefix: Option<&[u8]>, - mut f: F, - child_info: Option<&ChildInfo>, - maybe_start_at: Option<&[u8]>, - ) { - let mut iter = move |db| -> sp_std::result::Result<(), Box>> { - self.with_recorder_and_cache(Some(*root), |recorder, cache| { - let trie = TrieDBBuilder::::new(db, root) - .with_optional_recorder(recorder) - .with_optional_cache(cache) - .build(); - let prefix = maybe_prefix.unwrap_or(&[]); - let iter = match maybe_start_at { - Some(start_at) => - TrieDBKeyIterator::new_prefixed_then_seek(&trie, prefix, start_at), - None => TrieDBKeyIterator::new_prefixed(&trie, prefix), - }?; - - for x in iter { - let key = x?; - - debug_assert!(maybe_prefix - .as_ref() - .map(|prefix| key.starts_with(prefix)) - .unwrap_or(true)); - - if !f(&key) { - break - } - } - - Ok(()) - }) - }; - - let result = if let Some(child_info) = child_info { - let db = KeySpacedDB::new(self, child_info.keyspace()); - iter(&db) - } else { - iter(self) - }; - if let Err(e) = result { - debug!(target: "trie", "Error while iterating by prefix: {}", e); + if self.root == Default::default() { + // A special-case for an empty storage root. + return Ok(Default::default()) } - } - fn trie_iter_inner, Vec) -> bool>( - &self, - root: &H::Out, - prefix: Option<&[u8]>, - mut f: F, - child_info: Option<&ChildInfo>, - start_at: Option<&[u8]>, - allow_missing_nodes: bool, - ) -> Result { - let mut iter = move |db| -> sp_std::result::Result>> { - self.with_recorder_and_cache(Some(*root), |recorder, cache| { - let trie = TrieDBBuilder::::new(db, root) - .with_optional_recorder(recorder) - .with_optional_cache(cache) - .build(); - - let prefix = prefix.unwrap_or(&[]); - let iterator = if let Some(start_at) = start_at { - TrieDBIterator::new_prefixed_then_seek(&trie, prefix, start_at)? + let trie_iter = self + .with_trie_db(root, args.child_info.as_ref(), |db| { + let prefix = args.prefix.as_ref().map(|prefix| &**prefix).unwrap_or(&[]); + if let Some(start_at) = args.start_at { + TrieDBRawIterator::new_prefixed_then_seek(db, prefix, &start_at) } else { - TrieDBIterator::new_prefixed(&trie, prefix)? - }; - for x in iterator { - let (key, value) = x?; - - debug_assert!(key.starts_with(prefix)); - - if !f(key, value) { - return Ok(false) - } + TrieDBRawIterator::new_prefixed(db, prefix) } - - Ok(true) }) - }; + .map_err(|e| format!("TrieDB iteration error: {}", e))?; - let result = if let Some(child_info) = child_info { - let db = KeySpacedDB::new(self, child_info.keyspace()); - iter(&db) - } else { - iter(self) - }; - match result { - Ok(completed) => Ok(completed), - Err(e) if matches!(*e, TrieError::IncompleteDatabase(_)) && allow_missing_nodes => - Ok(false), - Err(e) => Err(format!("TrieDB iteration error: {}", e)), - } - } - - /// Execute given closure for all key and values starting with prefix. - pub fn for_key_values_with_prefix(&self, prefix: &[u8], mut f: F) { - let _ = self.trie_iter_inner( - &self.root, - Some(prefix), - |k, v| { - f(&k, &v); - true - }, - None, - None, - false, - ); - } - - /// Returns all `(key, value)` pairs in the trie. - pub fn pairs(&self) -> Vec<(StorageKey, StorageValue)> { - let collect_all = || -> sp_std::result::Result<_, Box>> { - self.with_recorder_and_cache(None, |recorder, cache| { - let trie = TrieDBBuilder::::new(self, self.root()) - .with_optional_cache(cache) - .with_optional_recorder(recorder) - .build(); - - let mut v = Vec::new(); - for x in trie.iter()? { - let (key, value) = x?; - v.push((key.to_vec(), value.to_vec())); - } - - Ok(v) - }) - }; - - match collect_all() { - Ok(v) => v, - Err(e) => { - debug!(target: "trie", "Error extracting trie values: {}", e); - Vec::new() - }, - } - } - - /// Returns all keys that start with the given `prefix`. - pub fn keys(&self, prefix: &[u8]) -> Vec { - let mut keys = Vec::new(); - self.for_keys_with_prefix(prefix, |k| keys.push(k.to_vec())); - keys + Ok(RawIter { + allow_missing: args.allow_missing, + child_info: args.child_info, + root, + trie_iter, + state: IterState::Pending, + _phantom: Default::default(), + }) } /// Return the storage root after applying the given `delta`. diff --git a/primitives/storage/src/lib.rs b/primitives/storage/src/lib.rs index 68d7b90639886..adb1fec6843aa 100644 --- a/primitives/storage/src/lib.rs +++ b/primitives/storage/src/lib.rs @@ -273,6 +273,7 @@ impl ChildInfo { /// Returns byte sequence (keyspace) that can be use by underlying db to isolate keys. /// This is a unique id of the child trie. The collision resistance of this value /// depends on the type of child info use. For `ChildInfo::Default` it is and need to be. + #[inline] pub fn keyspace(&self) -> &[u8] { match self { ChildInfo::ParentKeyId(..) => self.storage_key(), diff --git a/primitives/trie/Cargo.toml b/primitives/trie/Cargo.toml index 1e83a5d6dfec5..bc4f2aa777781 100644 --- a/primitives/trie/Cargo.toml +++ b/primitives/trie/Cargo.toml @@ -29,7 +29,7 @@ parking_lot = { version = "0.12.1", optional = true } scale-info = { version = "2.1.1", default-features = false, features = ["derive"] } thiserror = { version = "1.0.30", optional = true } tracing = { version = "0.1.29", optional = true } -trie-db = { version = "0.24.0", default-features = false } +trie-db = { version = "0.25.0", default-features = false } trie-root = { version = "0.17.0", default-features = false } sp-core = { version = "7.0.0", default-features = false, path = "../core" } sp-std = { version = "5.0.0", default-features = false, path = "../std" } diff --git a/primitives/trie/src/lib.rs b/primitives/trie/src/lib.rs index 01650e9a376be..716d37afbe9c1 100644 --- a/primitives/trie/src/lib.rs +++ b/primitives/trie/src/lib.rs @@ -51,7 +51,7 @@ pub use trie_db::{ nibble_ops, node::{NodePlan, ValuePlan}, CError, DBValue, Query, Recorder, Trie, TrieCache, TrieConfiguration, TrieDBIterator, - TrieDBKeyIterator, TrieLayout, TrieMut, TrieRecorder, + TrieDBKeyIterator, TrieDBRawIterator, TrieLayout, TrieMut, TrieRecorder, }; /// The Substrate format implementation of `TrieStream`. pub use trie_stream::TrieStream; @@ -442,6 +442,7 @@ fn keyspace_as_prefix_alloc(ks: &[u8], prefix: Prefix) -> (Vec, Option) impl<'a, DB: ?Sized, H> KeySpacedDB<'a, DB, H> { /// instantiate new keyspaced db + #[inline] pub fn new(db: &'a DB, ks: &'a [u8]) -> Self { KeySpacedDB(db, ks, PhantomData) } diff --git a/primitives/trie/src/recorder.rs b/primitives/trie/src/recorder.rs index bc67cfc287942..5d5b4d67d4791 100644 --- a/primitives/trie/src/recorder.rs +++ b/primitives/trie/src/recorder.rs @@ -83,6 +83,7 @@ impl Recorder { /// /// - `storage_root`: The storage root of the trie for which accesses are recorded. This is /// important when recording access to different tries at once (like top and child tries). + #[inline] pub fn as_trie_recorder( &self, storage_root: H::Out, diff --git a/test-utils/runtime/Cargo.toml b/test-utils/runtime/Cargo.toml index 2c943c47f32b0..e5987c75fef62 100644 --- a/test-utils/runtime/Cargo.toml +++ b/test-utils/runtime/Cargo.toml @@ -41,7 +41,7 @@ pallet-timestamp = { version = "4.0.0-dev", default-features = false, path = ".. sp-finality-grandpa = { version = "4.0.0-dev", default-features = false, path = "../../primitives/finality-grandpa" } sp-trie = { version = "7.0.0", default-features = false, path = "../../primitives/trie" } sp-transaction-pool = { version = "4.0.0-dev", default-features = false, path = "../../primitives/transaction-pool" } -trie-db = { version = "0.24.0", default-features = false } +trie-db = { version = "0.25.0", default-features = false } sc-service = { version = "0.10.0-dev", default-features = false, optional = true, features = ["test-helpers"], path = "../../client/service" } sp-state-machine = { version = "0.13.0", default-features = false, path = "../../primitives/state-machine" } sp-externalities = { version = "0.13.0", default-features = false, path = "../../primitives/externalities" } diff --git a/utils/frame/benchmarking-cli/src/storage/cmd.rs b/utils/frame/benchmarking-cli/src/storage/cmd.rs index ce2d52e57d641..c63b40dfa64f7 100644 --- a/utils/frame/benchmarking-cli/src/storage/cmd.rs +++ b/utils/frame/benchmarking-cli/src/storage/cmd.rs @@ -20,7 +20,6 @@ use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider}; use sc_client_db::DbHash; use sc_service::Configuration; use sp_blockchain::HeaderBackend; -use sp_core::storage::StorageKey; use sp_database::{ColumnId, Database}; use sp_runtime::traits::{Block as BlockT, HashFor}; use sp_state_machine::Storage; @@ -192,8 +191,7 @@ impl StorageCmd { BA: ClientBackend, { let hash = client.usage_info().chain.best_hash; - let empty_prefix = StorageKey(Vec::new()); - let mut keys = client.storage_keys(hash, &empty_prefix)?; + let mut keys: Vec<_> = client.storage_keys(hash, None, None)?.collect(); let (mut rng, _) = new_rng(None); keys.shuffle(&mut rng); diff --git a/utils/frame/benchmarking-cli/src/storage/read.rs b/utils/frame/benchmarking-cli/src/storage/read.rs index 20c41e4a5196b..501b0a43b84ba 100644 --- a/utils/frame/benchmarking-cli/src/storage/read.rs +++ b/utils/frame/benchmarking-cli/src/storage/read.rs @@ -17,7 +17,6 @@ use sc_cli::Result; use sc_client_api::{Backend as ClientBackend, StorageProvider, UsageProvider}; -use sp_core::storage::StorageKey; use sp_runtime::traits::{Block as BlockT, Header as HeaderT}; use log::info; @@ -42,8 +41,7 @@ impl StorageCmd { info!("Preparing keys from block {}", best_hash); // Load all keys and randomly shuffle them. - let empty_prefix = StorageKey(Vec::new()); - let mut keys = client.storage_keys(best_hash, &empty_prefix)?; + let mut keys: Vec<_> = client.storage_keys(best_hash, None, None)?.collect(); let (mut rng, _) = new_rng(None); keys.shuffle(&mut rng); @@ -55,8 +53,7 @@ impl StorageCmd { match (self.params.include_child_trees, self.is_child_key(key.clone().0)) { (true, Some(info)) => { // child tree key - let child_keys = client.child_storage_keys(best_hash, &info, &empty_prefix)?; - for ck in child_keys { + for ck in client.child_storage_keys(best_hash, info.clone(), None, None)? { child_nodes.push((ck.clone(), info.clone())); } }, diff --git a/utils/frame/benchmarking-cli/src/storage/write.rs b/utils/frame/benchmarking-cli/src/storage/write.rs index 1917bc5f8fd07..73c2438f22bc9 100644 --- a/utils/frame/benchmarking-cli/src/storage/write.rs +++ b/utils/frame/benchmarking-cli/src/storage/write.rs @@ -61,7 +61,10 @@ impl StorageCmd { info!("Preparing keys from block {}", best_hash); // Load all KV pairs and randomly shuffle them. - let mut kvs = trie.pairs(); + let mut kvs: Vec<_> = trie + .pairs(Default::default()) + .expect("iterating over storage shouldn't fail") + .collect(); let (mut rng, _) = new_rng(None); kvs.shuffle(&mut rng); info!("Writing {} keys", kvs.len()); @@ -70,11 +73,12 @@ impl StorageCmd { // 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 (k, original_v) in kvs { + for key_value in kvs { + let (k, original_v) = key_value.expect("iterating over storage shouldn't fail"); match (self.params.include_child_trees, self.is_child_key(k.to_vec())) { (true, Some(info)) => { let child_keys = - client.child_storage_keys_iter(best_hash, info.clone(), None, None)?; + client.child_storage_keys(best_hash, info.clone(), None, None)?; for ck in child_keys { child_nodes.push((ck.clone(), info.clone())); } diff --git a/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml b/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml index 17a68e2f4cbe8..c7c2f442b56f7 100644 --- a/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml +++ b/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml @@ -21,7 +21,7 @@ log = { version = "0.4.17", default-features = false } sp-core = { path = "../../../../primitives/core" } sp-state-machine = { path = "../../../../primitives/state-machine" } sp-trie = { path = "../../../../primitives/trie" } -trie-db = "0.24.0" +trie-db = "0.25.0" jsonrpsee = { version = "0.16.2", features = ["client-core", "server", "macros"] } From 8497910da7facc7ee297d3592abf55bce90cca79 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 31 Jan 2023 14:09:48 +0000 Subject: [PATCH 02/22] Make sure storage iteration is also accounted for when benchmarking --- Cargo.lock | 1 + client/db/Cargo.toml | 1 + client/db/src/bench.rs | 139 ++++++++++++++++++++++++++++++----------- 3 files changed, 106 insertions(+), 35 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2dafc0b5dbb26..fc2a8a4e66298 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8037,6 +8037,7 @@ dependencies = [ name = "sc-client-db" version = "0.10.0-dev" dependencies = [ + "array-bytes", "criterion", "hash-db", "kitchensink-runtime", diff --git a/client/db/Cargo.toml b/client/db/Cargo.toml index 7a62cb1ea0f53..73eca1dff708a 100644 --- a/client/db/Cargo.toml +++ b/client/db/Cargo.toml @@ -44,6 +44,7 @@ quickcheck = { version = "1.0.3", default-features = false } kitchensink-runtime = { path = "../../bin/node/runtime" } sp-tracing = { version = "6.0.0", path = "../../primitives/tracing" } substrate-test-runtime-client = { version = "2.0.0", path = "../../test-utils/runtime/client" } +array-bytes = "4.1" [features] default = [] diff --git a/client/db/src/bench.rs b/client/db/src/bench.rs index 965999d0690dc..30816eb37065a 100644 --- a/client/db/src/bench.rs +++ b/client/db/src/bench.rs @@ -22,6 +22,7 @@ use crate::{DbState, DbStateBuilder}; use hash_db::{Hasher, Prefix}; use kvdb::{DBTransaction, KeyValueDB}; use linked_hash_map::LinkedHashMap; +use parking_lot::Mutex; use sp_core::{ hexdisplay::HexDisplay, storage::{ChildInfo, TrackedStorageKey}, @@ -60,6 +61,19 @@ impl sp_state_machine::Storage> for StorageDb, TrackedStorageKey>, + /// Key tracker for keys in a child trie. + /// Child trie are identified by their storage key (i.e. `ChildInfo::storage_key()`) + /// We track the total number of reads and writes to these keys, + /// not de-duplicated for repeats. + child_keys: LinkedHashMap, LinkedHashMap, TrackedStorageKey>>, +} + /// State that manages the backend database reference. Allows runtime to control the database. pub struct BenchmarkingState { root: Cell, @@ -68,25 +82,18 @@ pub struct BenchmarkingState { db: Cell>>, genesis: HashMap, (Vec, i32)>, record: Cell>>, - /// Key tracker for keys in the main trie. - /// We track the total number of reads and writes to these keys, - /// not de-duplicated for repeats. - main_key_tracker: RefCell, TrackedStorageKey>>, - /// Key tracker for keys in a child trie. - /// Child trie are identified by their storage key (i.e. `ChildInfo::storage_key()`) - /// We track the total number of reads and writes to these keys, - /// not de-duplicated for repeats. - child_key_tracker: RefCell, LinkedHashMap, TrackedStorageKey>>>, + key_tracker: Arc>, whitelist: RefCell>, proof_recorder: Option>>, proof_recorder_root: Cell, - enable_tracking: bool, shared_trie_cache: SharedTrieCache>, } /// A raw iterator over the `BenchmarkingState`. pub struct RawIter { inner: as StateBackend>>::RawIter, + child_trie: Option>, + key_tracker: Arc>, } impl StorageIterator> for RawIter { @@ -94,14 +101,30 @@ impl StorageIterator> for RawIter { type Error = String; fn next_key(&mut self, backend: &Self::Backend) -> Option> { - self.inner.next_key(backend.state.borrow().as_ref()?) + match self.inner.next_key(backend.state.borrow().as_ref()?) { + Some(Ok(key)) => { + self.key_tracker + .lock() + .add_read_key(self.child_trie.as_ref().map(|child_trie| &**child_trie), &key); + Some(Ok(key)) + }, + result => result, + } } fn next_pair( &mut self, backend: &Self::Backend, ) -> Option> { - self.inner.next_pair(backend.state.borrow().as_ref()?) + match self.inner.next_pair(backend.state.borrow().as_ref()?) { + Some(Ok((key, value))) => { + self.key_tracker + .lock() + .add_read_key(self.child_trie.as_ref().map(|child_trie| &**child_trie), &key); + Some(Ok((key, value))) + }, + result => result, + } } fn was_complete(&self) -> bool { @@ -129,12 +152,14 @@ impl BenchmarkingState { genesis: Default::default(), genesis_root: Default::default(), record: Default::default(), - main_key_tracker: Default::default(), - child_key_tracker: Default::default(), + key_tracker: Arc::new(Mutex::new(KeyTracker { + main_keys: Default::default(), + child_keys: Default::default(), + enable_tracking, + })), whitelist: Default::default(), proof_recorder: record_proof.then(Default::default), proof_recorder_root: Cell::new(root), - enable_tracking, // Enable the cache, but do not sync anything to the shared state. shared_trie_cache: SharedTrieCache::new(CacheSize::new(0)), }; @@ -183,36 +208,51 @@ impl BenchmarkingState { } fn add_whitelist_to_tracker(&self) { - let mut main_key_tracker = self.main_key_tracker.borrow_mut(); + self.key_tracker.lock().add_whitelist(&self.whitelist.borrow()); + } - let whitelist = self.whitelist.borrow(); + fn wipe_tracker(&self) { + let mut key_tracker = self.key_tracker.lock(); + key_tracker.main_keys = LinkedHashMap::new(); + key_tracker.child_keys = LinkedHashMap::new(); + key_tracker.add_whitelist(&self.whitelist.borrow()); + } + fn add_read_key(&self, childtrie: Option<&[u8]>, key: &[u8]) { + self.key_tracker.lock().add_read_key(childtrie, key); + } + + fn add_write_key(&self, childtrie: Option<&[u8]>, key: &[u8]) { + self.key_tracker.lock().add_write_key(childtrie, key); + } + + fn all_trackers(&self) -> Vec { + self.key_tracker.lock().all_trackers() + } +} + +impl KeyTracker { + fn add_whitelist(&mut self, whitelist: &[TrackedStorageKey]) { whitelist.iter().for_each(|key| { let mut whitelisted = TrackedStorageKey::new(key.key.clone()); whitelisted.whitelist(); - main_key_tracker.insert(key.key.clone(), whitelisted); + self.main_keys.insert(key.key.clone(), whitelisted); }); } - fn wipe_tracker(&self) { - *self.main_key_tracker.borrow_mut() = LinkedHashMap::new(); - *self.child_key_tracker.borrow_mut() = LinkedHashMap::new(); - self.add_whitelist_to_tracker(); - } - // Childtrie is identified by its storage key (i.e. `ChildInfo::storage_key`) - fn add_read_key(&self, childtrie: Option<&[u8]>, key: &[u8]) { + fn add_read_key(&mut self, childtrie: Option<&[u8]>, key: &[u8]) { if !self.enable_tracking { return } - let mut child_key_tracker = self.child_key_tracker.borrow_mut(); - let mut main_key_tracker = self.main_key_tracker.borrow_mut(); + let child_key_tracker = &mut self.child_keys; + let main_key_tracker = &mut self.main_keys; let key_tracker = if let Some(childtrie) = childtrie { child_key_tracker.entry(childtrie.to_vec()).or_insert_with(LinkedHashMap::new) } else { - &mut main_key_tracker + main_key_tracker }; let should_log = match key_tracker.get_mut(key) { @@ -242,18 +282,18 @@ impl BenchmarkingState { } // Childtrie is identified by its storage key (i.e. `ChildInfo::storage_key`) - fn add_write_key(&self, childtrie: Option<&[u8]>, key: &[u8]) { + fn add_write_key(&mut self, childtrie: Option<&[u8]>, key: &[u8]) { if !self.enable_tracking { return } - let mut child_key_tracker = self.child_key_tracker.borrow_mut(); - let mut main_key_tracker = self.main_key_tracker.borrow_mut(); + let child_key_tracker = &mut self.child_keys; + let main_key_tracker = &mut self.main_keys; let key_tracker = if let Some(childtrie) = childtrie { child_key_tracker.entry(childtrie.to_vec()).or_insert_with(LinkedHashMap::new) } else { - &mut main_key_tracker + main_key_tracker }; // If we have written to the key, we also consider that we have read from it. @@ -287,11 +327,11 @@ impl BenchmarkingState { fn all_trackers(&self) -> Vec { let mut all_trackers = Vec::new(); - self.main_key_tracker.borrow().iter().for_each(|(_, tracker)| { + self.main_keys.iter().for_each(|(_, tracker)| { all_trackers.push(tracker.clone()); }); - self.child_key_tracker.borrow().iter().for_each(|(_, child_tracker)| { + self.child_keys.iter().for_each(|(_, child_tracker)| { child_tracker.iter().for_each(|(_, tracker)| { all_trackers.push(tracker.clone()); }); @@ -413,12 +453,18 @@ impl StateBackend> for BenchmarkingState { } fn raw_iter(&self, args: IterArgs) -> Result { + let child_trie = + args.child_info.as_ref().map(|child_info| child_info.storage_key().to_vec()); self.state .borrow() .as_ref() .map(|s| s.raw_iter(args)) .unwrap_or(Ok(Default::default())) - .map(|raw_iter| RawIter { inner: raw_iter }) + .map(|raw_iter| RawIter { + inner: raw_iter, + key_tracker: self.key_tracker.clone(), + child_trie, + }) } fn commit( @@ -608,6 +654,29 @@ mod test { use crate::bench::BenchmarkingState; use sp_state_machine::backend::Backend as _; + fn hex(hex: &str) -> Vec { + array_bytes::hex2bytes(hex).unwrap() + } + + #[test] + fn iteration_is_also_counted_in_rw_counts() { + let storage = sp_runtime::Storage { + top: vec![( + hex("ce6e1397e668c7fcf47744350dc59688455a2c2dbd2e2a649df4e55d93cd7158"), + hex("0102030405060708"), + )] + .into_iter() + .collect(), + ..sp_runtime::Storage::default() + }; + let bench_state = + BenchmarkingState::::new(storage, None, false, true).unwrap(); + + assert_eq!(bench_state.read_write_count(), (0, 0, 0, 0)); + assert_eq!(bench_state.keys(Default::default()).unwrap().count(), 1); + assert_eq!(bench_state.read_write_count(), (1, 0, 0, 0)); + } + #[test] fn read_to_main_and_child_tries() { let bench_state = From e6b35c9bc442a11b614abfce3bab0b28fc99aa18 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 09:38:30 +0000 Subject: [PATCH 03/22] Use `trie-db` from crates.io --- Cargo.lock | 6 ++++-- Cargo.toml | 4 ---- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fc2a8a4e66298..e9e05c8b2251a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2847,7 +2847,8 @@ dependencies = [ [[package]] name = "hash-db" version = "0.15.2" -source = "git+https://github.com/koute/trie.git?branch=master_iterators#a815264f668deb92a06c27233f5e899a60790049" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d23bd4e7b5eda0d0f3a307e8b381fdc8ba9000f26fbe912250c0a4cc3956364a" [[package]] name = "hash256-std-hasher" @@ -11346,7 +11347,8 @@ dependencies = [ [[package]] name = "trie-db" version = "0.25.0" -source = "git+https://github.com/koute/trie.git?branch=master_iterators#4d8bda6078eb911027ee3f43440a4272ab673ee5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bcd4e7fbb3bcab17b02596b53876e36eb39663ddd87884bfd88c69cdeb2ebd6" dependencies = [ "hash-db", "hashbrown 0.13.2", diff --git a/Cargo.toml b/Cargo.toml index a222a050c6107..b004886e83f9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -321,7 +321,3 @@ inherits = "release" lto = "fat" # https://doc.rust-lang.org/rustc/codegen-options/index.html#codegen-units codegen-units = 1 - -[patch.crates-io] -trie-db = { git = "https://github.com/koute/trie.git", branch = "master_iterators" } -hash-db = { git = "https://github.com/koute/trie.git", branch = "master_iterators" } From 255658142720227d319cd852bcce478aa3803c41 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 09:48:14 +0000 Subject: [PATCH 04/22] Appease clippy --- primitives/state-machine/src/trie_backend_essence.rs | 2 +- primitives/trie/src/cache/mod.rs | 4 ++-- primitives/trie/src/recorder.rs | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 7c12c86127e96..297bd49382f0b 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -606,7 +606,7 @@ where let trie_iter = self .with_trie_db(root, args.child_info.as_ref(), |db| { - let prefix = args.prefix.as_ref().map(|prefix| &**prefix).unwrap_or(&[]); + let prefix = args.prefix.as_deref().unwrap_or(&[]); if let Some(start_at) = args.start_at { TrieDBRawIterator::new_prefixed_then_seek(db, prefix, &start_at) } else { diff --git a/primitives/trie/src/cache/mod.rs b/primitives/trie/src/cache/mod.rs index 3c1e5b8d0ff0b..332dc8f470838 100644 --- a/primitives/trie/src/cache/mod.rs +++ b/primitives/trie/src/cache/mod.rs @@ -438,8 +438,8 @@ enum ValueCache<'a, H: Hasher> { impl ValueCache<'_, H> { /// Get the value for the given `key`. - fn get<'a>( - &'a mut self, + fn get( + &mut self, key: &[u8], shared_cache: &SharedTrieCache, stats: &HitStats, diff --git a/primitives/trie/src/recorder.rs b/primitives/trie/src/recorder.rs index 5d5b4d67d4791..7d19ec00b16d9 100644 --- a/primitives/trie/src/recorder.rs +++ b/primitives/trie/src/recorder.rs @@ -148,7 +148,7 @@ struct TrieRecorder { impl>> trie_db::TrieRecorder for TrieRecorder { - fn record<'b>(&mut self, access: TrieAccess<'b, H::Out>) { + fn record(&mut self, access: TrieAccess) { let mut encoded_size_update = 0; match access { From 1f1ff7a0dfc40983d491587315fabe79fc621531 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 11:32:14 +0000 Subject: [PATCH 05/22] Bump `trie-bench` to 0.35.0 --- Cargo.lock | 27 +++++++-------------------- primitives/trie/Cargo.toml | 2 +- 2 files changed, 8 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 082d0544f27a1..779aca6c65f3f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10259,7 +10259,7 @@ dependencies = [ "sp-trie", "thiserror", "tracing", - "trie-db 0.25.0", + "trie-db", ] [[package]] @@ -10359,7 +10359,7 @@ dependencies = [ "thiserror", "tracing", "trie-bench", - "trie-db 0.25.0", + "trie-db", "trie-root", "trie-standardmap", ] @@ -10675,7 +10675,7 @@ dependencies = [ "sp-runtime", "sp-state-machine", "sp-trie", - "trie-db 0.25.0", + "trie-db", ] [[package]] @@ -10747,7 +10747,7 @@ dependencies = [ "sp-version", "substrate-test-runtime-client", "substrate-wasm-builder", - "trie-db 0.25.0", + "trie-db", ] [[package]] @@ -11321,33 +11321,20 @@ dependencies = [ [[package]] name = "trie-bench" -version = "0.34.0" +version = "0.35.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fbb0a830db7c42ae97ce4e21b30e2cf9dbcc1b4f7853bd1aedad3d806c281d0" +checksum = "22c1d18c423077531e693e87ace54ed9e4af1e4ce0a3ea8c9aa6608741074e2b" dependencies = [ "criterion", "hash-db", "keccak-hasher", "memory-db", "parity-scale-codec", - "trie-db 0.24.0", + "trie-db", "trie-root", "trie-standardmap", ] -[[package]] -name = "trie-db" -version = "0.24.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "004e1e8f92535694b4cb1444dc5a8073ecf0815e3357f729638b9f8fc4062908" -dependencies = [ - "hash-db", - "hashbrown 0.12.3", - "log", - "rustc-hex", - "smallvec", -] - [[package]] name = "trie-db" version = "0.25.0" diff --git a/primitives/trie/Cargo.toml b/primitives/trie/Cargo.toml index bc4f2aa777781..33a62cdd94084 100644 --- a/primitives/trie/Cargo.toml +++ b/primitives/trie/Cargo.toml @@ -38,7 +38,7 @@ schnellru = { version = "0.2.1", optional = true } [dev-dependencies] array-bytes = "4.1" criterion = "0.4.0" -trie-bench = "0.34.0" +trie-bench = "0.35.0" trie-standardmap = "0.15.2" sp-runtime = { version = "7.0.0", path = "../runtime" } From fa3edbd8e257b658ab68caa19d307d9d22064325 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 11:55:53 +0000 Subject: [PATCH 06/22] Fix tests' compilation --- client/service/test/src/client/mod.rs | 29 ++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/client/service/test/src/client/mod.rs b/client/service/test/src/client/mod.rs index 4b84d1496256e..af96a02f72e04 100644 --- a/client/service/test/src/client/mod.rs +++ b/client/service/test/src/client/mod.rs @@ -347,7 +347,20 @@ fn block_builder_works_with_transactions() { .expect("block 1 was just imported. qed"); assert_eq!(client.chain_info().best_number, 1); - assert_ne!(client.state_at(hash1).unwrap().pairs(), client.state_at(hash0).unwrap().pairs()); + assert_ne!( + client + .state_at(hash1) + .unwrap() + .pairs(Default::default()) + .unwrap() + .collect::>(), + client + .state_at(hash0) + .unwrap() + .pairs(Default::default()) + .unwrap() + .collect::>() + ); assert_eq!( client .runtime_api() @@ -406,8 +419,18 @@ fn block_builder_does_not_include_invalid() { assert_eq!(client.chain_info().best_number, 1); assert_ne!( - client.state_at(hashof1).unwrap().pairs(), - client.state_at(hashof0).unwrap().pairs() + client + .state_at(hashof1) + .unwrap() + .pairs(Default::default()) + .unwrap() + .collect::>(), + client + .state_at(hashof0) + .unwrap() + .pairs(Default::default()) + .unwrap() + .collect::>() ); assert_eq!(client.body(hashof1).unwrap().unwrap().len(), 1) } From 4be0cadea1d0a1feb17f4891a0fa322ca9880609 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:03:19 +0000 Subject: [PATCH 07/22] Update comment to clarify how `IterArgs::start_at` works --- primitives/state-machine/src/backend.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/primitives/state-machine/src/backend.rs b/primitives/state-machine/src/backend.rs index 551090fb1e228..d8bc2ce375b89 100644 --- a/primitives/state-machine/src/backend.rs +++ b/primitives/state-machine/src/backend.rs @@ -38,7 +38,9 @@ pub struct IterArgs<'a> { /// The prefix of the keys over which to iterate. pub prefix: Option<&'a [u8]>, - /// The key from which to start the iteration from. + /// The prefix from which to start the iteration from. + /// + /// This is inclusive and the iteration will include the key which is specified here. pub start_at: Option<&'a [u8]>, /// The info of the child trie over which to iterate over. From 825ba75e8b143f75824ed219bfc04e0250a505ac Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:04:01 +0000 Subject: [PATCH 08/22] Add extra tests --- primitives/state-machine/src/trie_backend.rs | 151 +++++++++++++++++++ 1 file changed, 151 insertions(+) diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index 312933716b675..29b9ffda1259d 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -559,6 +559,157 @@ pub mod tests { .is_none()); } + parameterized_test!(storage_iteration_works, storage_iteration_works_inner); + fn storage_iteration_works_inner( + state_version: StateVersion, + cache: Option, + recorder: Option, + ) { + let trie = test_trie(state_version, cache, recorder); + + // Fetch everything. + assert_eq!( + trie.keys(Default::default()) + .unwrap() + .map(|result| result.unwrap()) + .take(5) + .collect::>(), + vec![ + b":child_storage:default:sub1".to_vec(), + b":code".to_vec(), + b"key".to_vec(), + b"value1".to_vec(), + b"value2".to_vec(), + ] + ); + + // Fetch starting at a given key (full key). + assert_eq!( + trie.keys(IterArgs { start_at: Some(b"key"), ..IterArgs::default() }) + .unwrap() + .map(|result| result.unwrap()) + .take(3) + .collect::>(), + vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),] + ); + + // Fetch starting at a given key (partial key). + assert_eq!( + trie.keys(IterArgs { start_at: Some(b"ke"), ..IterArgs::default() }) + .unwrap() + .map(|result| result.unwrap()) + .take(3) + .collect::>(), + vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),] + ); + + // Fetch starting at a given key (empty key). + assert_eq!( + trie.keys(IterArgs { start_at: Some(b""), ..IterArgs::default() }) + .unwrap() + .map(|result| result.unwrap()) + .take(5) + .collect::>(), + vec![ + b":child_storage:default:sub1".to_vec(), + b":code".to_vec(), + b"key".to_vec(), + b"value1".to_vec(), + b"value2".to_vec(), + ] + ); + + // Fetch starting at a given key and with prefix which doesn't match that key. + assert!(trie + .keys(IterArgs { + prefix: Some(b"value"), + start_at: Some(b"key"), + ..IterArgs::default() + }) + .unwrap() + .map(|result| result.unwrap()) + .next() + .is_none()); + + // Fetch starting at a given key and with prefix which does match that key. + assert_eq!( + trie.keys(IterArgs { + prefix: Some(b"value"), + start_at: Some(b"value"), + ..IterArgs::default() + }) + .unwrap() + .map(|result| result.unwrap()) + .collect::>(), + vec![b"value1".to_vec(), b"value2".to_vec(),] + ); + + // Also test out the wrapper methods. + // TODO: Remove this once these methods are gone. + + let mut list = Vec::new(); + assert!(trie + .apply_to_key_values_while( + None, + None, + Some(b"key"), + |key, _| { + list.push(key); + true + }, + false + ) + .unwrap()); + assert_eq!(list[0..3], vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),]); + + let mut list = Vec::new(); + trie.apply_to_keys_while(None, None, Some(b"key"), |key| { + list.push(key.to_vec()); + true + }) + .unwrap(); + assert_eq!(list[0..3], vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),]); + + let mut list = Vec::new(); + trie.apply_to_keys_while(None, None, Some(b"k"), |key| { + list.push(key.to_vec()); + true + }); + assert_eq!(list[0..3], vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),]); + + let mut list = Vec::new(); + trie.apply_to_keys_while(None, None, Some(b""), |key| { + list.push(key.to_vec()); + true + }); + assert_eq!( + list[0..5], + vec![ + b":child_storage:default:sub1".to_vec(), + b":code".to_vec(), + b"key".to_vec(), + b"value1".to_vec(), + b"value2".to_vec(), + ] + ); + + let mut list = Vec::new(); + trie.apply_to_keys_while(None, Some(b"value"), Some(b"key"), |key| { + list.push(key.to_vec()); + true + }) + .unwrap(); + assert!(list.is_empty()); + + let mut list = Vec::new(); + trie.apply_to_keys_while(None, Some(b"value"), Some(b"value"), |key| { + list.push(key.to_vec()); + true + }) + .unwrap(); + assert_eq!(list, vec![b"value1".to_vec(), b"value2".to_vec(),]); + } + parameterized_test!(storage_root_is_non_default, storage_root_is_non_default_inner); fn storage_root_is_non_default_inner( state_version: StateVersion, From 65d99edfd3dd2522ca414b21d066da51be258124 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:05:19 +0000 Subject: [PATCH 09/22] Fix iterators on `Client` so that they behave as before --- client/api/src/backend.rs | 46 +++++++++++++++++++++---- client/service/test/src/client/mod.rs | 48 +++++++++++++++++++++++++++ 2 files changed, 88 insertions(+), 6 deletions(-) diff --git a/client/api/src/backend.rs b/client/api/src/backend.rs index 3880755f9908d..627fc4aa097f3 100644 --- a/client/api/src/backend.rs +++ b/client/api/src/backend.rs @@ -309,6 +309,7 @@ where { inner: >>::RawIter, state: State, + skip_if_first: Option, } impl KeysIter @@ -326,7 +327,12 @@ where args.prefix = prefix.as_ref().map(|prefix| prefix.0.as_slice()); args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); - Ok(Self { inner: state.raw_iter(args)?, state }) + let start_at = args.start_at; + Ok(Self { + inner: state.raw_iter(args)?, + state, + skip_if_first: start_at.map(|key| StorageKey(key.to_vec())), + }) } /// Create a new iterator over a child storage's keys. @@ -341,7 +347,12 @@ where args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); args.child_info = Some(child_info); - Ok(Self { inner: state.raw_iter(args)?, state }) + let start_at = args.start_at; + Ok(Self { + inner: state.raw_iter(args)?, + state, + skip_if_first: start_at.map(|key| StorageKey(key.to_vec())), + }) } } @@ -353,7 +364,15 @@ where type Item = StorageKey; fn next(&mut self) -> Option { - self.inner.next_key(&self.state)?.ok().map(StorageKey) + let key = self.inner.next_key(&self.state)?.ok().map(StorageKey)?; + + if let Some(skipped_key) = self.skip_if_first.take() { + if key == skipped_key { + return self.next() + } + } + + Some(key) } } @@ -365,6 +384,7 @@ where { inner: >>::RawIter, state: State, + skip_if_first: Option, } impl Iterator for PairsIter @@ -375,10 +395,19 @@ where type Item = (StorageKey, StorageData); fn next(&mut self) -> Option { - self.inner + let (key, value) = self + .inner .next_pair(&self.state)? .ok() - .map(|(key, value)| (StorageKey(key), StorageData(value))) + .map(|(key, value)| (StorageKey(key), StorageData(value)))?; + + if let Some(skipped_key) = self.skip_if_first.take() { + if key == skipped_key { + return self.next() + } + } + + Some((key, value)) } } @@ -397,7 +426,12 @@ where args.prefix = prefix.as_ref().map(|prefix| prefix.0.as_slice()); args.start_at = start_at.as_ref().map(|start_at| start_at.0.as_slice()); - Ok(Self { inner: state.raw_iter(args)?, state }) + let start_at = args.start_at; + Ok(Self { + inner: state.raw_iter(args)?, + state, + skip_if_first: start_at.map(|key| StorageKey(key.to_vec())), + }) } } diff --git a/client/service/test/src/client/mod.rs b/client/service/test/src/client/mod.rs index af96a02f72e04..872cd38950e63 100644 --- a/client/service/test/src/client/mod.rs +++ b/client/service/test/src/client/mod.rs @@ -1711,6 +1711,54 @@ fn storage_keys_works() { ] ); + // Starting at an empty key nothing gets skipped. + let res: Vec<_> = client + .storage_keys(block_hash, Some(&prefix), Some(&StorageKey("".into()))) + .unwrap() + .take(9) + .map(|x| array_bytes::bytes2hex("", &x.0)) + .collect(); + assert_eq!( + res, + [ + "00c232cf4e70a5e343317016dc805bf80a6a8cd8ad39958d56f99891b07851e0", + "085b2407916e53a86efeb8b72dbe338c4b341dab135252f96b6ed8022209b6cb", + "0befda6e1ca4ef40219d588a727f1271", + "1a560ecfd2a62c2b8521ef149d0804eb621050e3988ed97dca55f0d7c3e6aa34", + "1d66850d32002979d67dd29dc583af5b2ae2a1f71c1f35ad90fff122be7a3824", + "237498b98d8803334286e9f0483ef513098dd3c1c22ca21c4dc155b4ef6cc204", + "26aa394eea5630e07c48ae0c9558cef75e0621c4869aa60c02be9adcc98a0d1d", + "29b9db10ec5bf7907d8f74b5e60aa8140c4fbdd8127a1ee5600cb98e5ec01729", + "3a636f6465", + ] + ); + + // Starting at an incomplete key nothing gets skipped. + let res: Vec<_> = client + .storage_keys( + block_hash, + Some(&prefix), + Some(&StorageKey(array_bytes::hex2bytes_unchecked("3a636f64"))), + ) + .unwrap() + .take(8) + .map(|x| array_bytes::bytes2hex("", &x.0)) + .collect(); + assert_eq!( + res, + [ + "3a636f6465", + "3a686561707061676573", + "52008686cc27f6e5ed83a216929942f8bcd32a396f09664a5698f81371934b56", + "5348d72ac6cc66e5d8cbecc27b0e0677503b845fe2382d819f83001781788fd5", + "5c2d5fda66373dabf970e4fb13d277ce91c5233473321129d32b5a8085fa8133", + "6644b9b8bc315888ac8e41a7968dc2b4141a5403c58acdf70b7e8f7e07bf5081", + "66484000ed3f75c95fc7b03f39c20ca1e1011e5999278247d3b2f5e3c3273808", + "7d5007603a7f5dd729d51d93cf695d6465789443bb967c0d1fe270e388c96eaa", + ] + ); + + // Starting at a complete key the first key is skipped. let res: Vec<_> = client .storage_keys( block_hash, From ab0fff6c2d3ab44b8a3c15f2fd9297cd95126bde Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:20:59 +0000 Subject: [PATCH 10/22] Add extra `unwrap`s in tests --- primitives/state-machine/src/trie_backend.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index 29b9ffda1259d..9834b334562af 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -674,14 +674,16 @@ pub mod tests { trie.apply_to_keys_while(None, None, Some(b"k"), |key| { list.push(key.to_vec()); true - }); + }) + .unwrap(); assert_eq!(list[0..3], vec![b"key".to_vec(), b"value1".to_vec(), b"value2".to_vec(),]); let mut list = Vec::new(); trie.apply_to_keys_while(None, None, Some(b""), |key| { list.push(key.to_vec()); true - }); + }) + .unwrap(); assert_eq!( list[0..5], vec![ From 2e4160486e9f5cb383e36b089bf95b67e2f9fa9a Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:30:43 +0000 Subject: [PATCH 11/22] More clippy fixes --- client/db/src/bench.rs | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/client/db/src/bench.rs b/client/db/src/bench.rs index 30816eb37065a..537ca84bc52fe 100644 --- a/client/db/src/bench.rs +++ b/client/db/src/bench.rs @@ -103,9 +103,7 @@ impl StorageIterator> for RawIter { fn next_key(&mut self, backend: &Self::Backend) -> Option> { match self.inner.next_key(backend.state.borrow().as_ref()?) { Some(Ok(key)) => { - self.key_tracker - .lock() - .add_read_key(self.child_trie.as_ref().map(|child_trie| &**child_trie), &key); + self.key_tracker.lock().add_read_key(self.child_trie.as_deref(), &key); Some(Ok(key)) }, result => result, @@ -118,9 +116,7 @@ impl StorageIterator> for RawIter { ) -> Option> { match self.inner.next_pair(backend.state.borrow().as_ref()?) { Some(Ok((key, value))) => { - self.key_tracker - .lock() - .add_read_key(self.child_trie.as_ref().map(|child_trie| &**child_trie), &key); + self.key_tracker.lock().add_read_key(self.child_trie.as_deref(), &key); Some(Ok((key, value))) }, result => result, From b5ce5f59f0bd0c0f941d9a2a739fe3d51b681eb3 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Fri, 3 Feb 2023 14:50:50 +0000 Subject: [PATCH 12/22] Come on clippy, give me a break already --- client/service/src/client/client.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/client/service/src/client/client.rs b/client/service/src/client/client.rs index 01db050abfc9c..73d66b6fc041f 100644 --- a/client/service/src/client/client.rs +++ b/client/service/src/client/client.rs @@ -1471,8 +1471,8 @@ where start_key: Option<&StorageKey>, ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - Ok(KeysIter::new(state, prefix, start_key) - .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) + KeysIter::new(state, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e))) } fn child_storage_keys( @@ -1483,8 +1483,8 @@ where start_key: Option<&StorageKey>, ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - Ok(KeysIter::new_child(state, child_info, prefix, start_key) - .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) + KeysIter::new_child(state, child_info, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e))) } fn storage_pairs( @@ -1494,8 +1494,8 @@ where start_key: Option<&StorageKey>, ) -> sp_blockchain::Result> { let state = self.state_at(hash)?; - Ok(PairsIter::new(state, prefix, start_key) - .map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))?) + PairsIter::new(state, prefix, start_key) + .map_err(|e| sp_blockchain::Error::from_state(Box::new(e))) } fn storage( From e04904244faf5e3566e4d18407449bf0be5c45bf Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 7 Feb 2023 09:17:56 +0000 Subject: [PATCH 13/22] Rename `allow_missing` to `stop_on_incomplete_database` --- primitives/state-machine/src/backend.rs | 4 ++-- primitives/state-machine/src/trie_backend_essence.rs | 10 ++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/primitives/state-machine/src/backend.rs b/primitives/state-machine/src/backend.rs index d8bc2ce375b89..1c098eb6e8602 100644 --- a/primitives/state-machine/src/backend.rs +++ b/primitives/state-machine/src/backend.rs @@ -51,7 +51,7 @@ pub struct IterArgs<'a> { /// When a missing trie node is reached the iterator will: /// - return an error if this is set to `false` (default) /// - return `None` if this is set to `true` - pub allow_missing: bool, + pub stop_on_incomplete_database: bool, } /// A trait for a raw storage iterator. @@ -236,7 +236,7 @@ pub trait Backend: sp_std::fmt::Debug { child_info: child_info.cloned(), prefix, start_at, - allow_missing, + stop_on_incomplete_database: allow_missing, ..IterArgs::default() }; let mut iter = self.pairs(args)?; diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 297bd49382f0b..3a41d88fac042 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -266,7 +266,7 @@ pub struct RawIter where H: Hasher, { - allow_missing: bool, + stop_on_incomplete_database: bool, root: H::Out, child_info: Option, trie_iter: TrieDBRawIterator>, @@ -305,7 +305,9 @@ where }, Some(Err(error)) => { self.state = IterState::FinishedIncomplete; - if matches!(*error, TrieError::IncompleteDatabase(_)) && self.allow_missing { + if matches!(*error, TrieError::IncompleteDatabase(_)) && + self.stop_on_incomplete_database + { None } else { Some(Err(format!("TrieDB iteration error: {}", error))) @@ -321,7 +323,7 @@ where { fn default() -> Self { Self { - allow_missing: false, + stop_on_incomplete_database: false, child_info: None, root: Default::default(), trie_iter: TrieDBRawIterator::empty(), @@ -616,7 +618,7 @@ where .map_err(|e| format!("TrieDB iteration error: {}", e))?; Ok(RawIter { - allow_missing: args.allow_missing, + stop_on_incomplete_database: args.stop_on_incomplete_database, child_info: args.child_info, root, trie_iter, From 387c600a80b07fad5776582c3028193cd50836c7 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 07:25:12 +0000 Subject: [PATCH 14/22] Add `#[inline]` to `with_recorder_and_cache` --- primitives/state-machine/src/trie_backend_essence.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 3a41d88fac042..8647ea6ddf868 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -167,6 +167,7 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss /// /// If the given `storage_root` is `None`, `self.root` will be used. #[cfg(feature = "std")] + #[inline] fn with_recorder_and_cache( &self, storage_root: Option, @@ -192,6 +193,7 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss } #[cfg(not(feature = "std"))] + #[inline] fn with_recorder_and_cache( &self, _: Option, From 962cdbc82749570338861ecccbf9e7c2907d34bd Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 07:26:03 +0000 Subject: [PATCH 15/22] Use `with_recorder_and_cache` in `with_trie_db`; add doc comment --- .../state-machine/src/trie_backend_essence.rs | 85 ++++++------------- 1 file changed, 25 insertions(+), 60 deletions(-) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 8647ea6ddf868..28409d5fb4a69 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -335,66 +335,6 @@ where } } -impl TrieBackendEssence -where - H: Hasher, - H::Out: Codec + Ord, - S: TrieBackendStorage, - C: AsLocalTrieCache + Send + Sync, -{ - #[cfg(not(feature = "std"))] - #[inline] - fn with_trie_db( - &self, - root: H::Out, - child_info: Option<&ChildInfo>, - callback: impl FnOnce(&sp_trie::TrieDB>) -> R, - ) -> R { - let backend = self as &dyn HashDBRef>; - let db = child_info - .as_ref() - .map(|child_info| KeySpacedDB::new(backend, child_info.keyspace())); - let db = db.as_ref().map(|db| db as &dyn HashDBRef>).unwrap_or(backend); - - let trie = TrieDBBuilder::::new(db, &root).build(); - - callback(&trie) - } - - #[cfg(feature = "std")] - #[inline] - fn with_trie_db( - &self, - root: H::Out, - child_info: Option<&ChildInfo>, - callback: impl FnOnce(&sp_trie::TrieDB>) -> R, - ) -> R { - let backend = self as &dyn HashDBRef>; - let db = child_info - .as_ref() - .map(|child_info| KeySpacedDB::new(backend, child_info.keyspace())); - let db = db.as_ref().map(|db| db as &dyn HashDBRef>).unwrap_or(backend); - - let recorder = self.recorder.as_ref(); - let mut recorder = recorder.as_ref().map(|r| r.as_trie_recorder(root)); - let recorder = match recorder.as_mut() { - Some(recorder) => Some(recorder as &mut dyn TrieRecorder), - None => None, - }; - - let cache = self.trie_node_cache.as_ref(); - let mut cache = cache.map(|c| c.as_local_trie_cache().as_trie_db_cache(root)); - let cache = cache.as_mut().map(|c| c as &mut dyn sp_trie::TrieCache>); - - let trie = TrieDBBuilder::::new(db, &root) - .with_optional_recorder(recorder) - .with_optional_cache(cache) - .build(); - - callback(&trie) - } -} - impl StorageIterator for RawIter where H: Hasher, @@ -425,6 +365,31 @@ impl, H: Hasher, C: AsLocalTrieCache + Send + Sync> where H::Out: Codec + Ord, { + /// Calls the given closure with a [`TrieDb`] constructed for the given + /// storage root and (optionally) child trie. + #[inline] + fn with_trie_db( + &self, + root: H::Out, + child_info: Option<&ChildInfo>, + callback: impl FnOnce(&sp_trie::TrieDB>) -> R, + ) -> R { + let backend = self as &dyn HashDBRef>; + let db = child_info + .as_ref() + .map(|child_info| KeySpacedDB::new(backend, child_info.keyspace())); + let db = db.as_ref().map(|db| db as &dyn HashDBRef>).unwrap_or(backend); + + self.with_recorder_and_cache(Some(root), |recorder, cache| { + let trie = TrieDBBuilder::::new(db, &root) + .with_optional_recorder(recorder) + .with_optional_cache(cache) + .build(); + + callback(&trie) + }) + } + /// Return the next key in the trie i.e. the minimum key that is strictly superior to `key` in /// lexicographic order. pub fn next_storage_key(&self, key: &[u8]) -> Result> { From 6582bc9be03cc2ccc8c4a854b7e680bc903aae94 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 07:30:49 +0000 Subject: [PATCH 16/22] Simplify code: use `with_trie_db` in `next_storage_key_from_root` --- .../state-machine/src/trie_backend_essence.rs | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 28409d5fb4a69..3788a58c0fe1b 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -444,21 +444,7 @@ where child_info: Option<&ChildInfo>, key: &[u8], ) -> Result> { - let dyn_eph: &dyn HashDBRef<_, _>; - let keyspace_eph; - if let Some(child_info) = child_info.as_ref() { - keyspace_eph = KeySpacedDB::new(self, child_info.keyspace()); - dyn_eph = &keyspace_eph; - } else { - dyn_eph = self; - } - - self.with_recorder_and_cache(Some(*root), |recorder, cache| { - let trie = TrieDBBuilder::::new(dyn_eph, root) - .with_optional_recorder(recorder) - .with_optional_cache(cache) - .build(); - + self.with_trie_db(*root, child_info, |trie| { let mut iter = trie.key_iter().map_err(|e| format!("TrieDB iteration error: {}", e))?; // The key just after the one given in input, basically `key++0`. From 0569cb83b0698815f966add4ca1d832c6380af5d Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 08:19:02 +0000 Subject: [PATCH 17/22] Remove `expect`s in the benchmarking CLI --- utils/frame/benchmarking-cli/src/storage/write.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/utils/frame/benchmarking-cli/src/storage/write.rs b/utils/frame/benchmarking-cli/src/storage/write.rs index 73c2438f22bc9..2a054c2805749 100644 --- a/utils/frame/benchmarking-cli/src/storage/write.rs +++ b/utils/frame/benchmarking-cli/src/storage/write.rs @@ -61,10 +61,7 @@ impl StorageCmd { info!("Preparing keys from block {}", best_hash); // Load all KV pairs and randomly shuffle them. - let mut kvs: Vec<_> = trie - .pairs(Default::default()) - .expect("iterating over storage shouldn't fail") - .collect(); + let mut kvs: Vec<_> = trie.pairs(Default::default())?.collect(); let (mut rng, _) = new_rng(None); kvs.shuffle(&mut rng); info!("Writing {} keys", kvs.len()); @@ -74,7 +71,7 @@ impl StorageCmd { // 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.expect("iterating over storage shouldn't fail"); + 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 = From 0c332c66c4f1f91073bae6cdd2321e7e3d2c78a7 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 08:21:36 +0000 Subject: [PATCH 18/22] Add extra doc comments --- primitives/state-machine/src/trie_backend_essence.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index 3788a58c0fe1b..e873b5e35c6e2 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -264,6 +264,7 @@ enum IterState { FinishedIncomplete, } +/// A raw iterator over the storage. pub struct RawIter where H: Hasher, @@ -543,6 +544,7 @@ where }) } + /// Create a raw iterator over the storage. pub fn raw_iter(&self, args: IterArgs) -> Result> { let root = if let Some(child_info) = args.child_info.as_ref() { let root = match self.child_root(&child_info)? { From 4717063b6bc4cc36fd8072f078aa7cfa96097516 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 08:22:40 +0000 Subject: [PATCH 19/22] Move `RawIter` before `TrieBackendEssence` (no code changes; just cut-paste) --- .../state-machine/src/trie_backend_essence.rs | 206 +++++++++--------- 1 file changed, 103 insertions(+), 103 deletions(-) diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index e873b5e35c6e2..dc49894e5b5e7 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -75,6 +75,109 @@ impl Cache { } } +enum IterState { + Pending, + FinishedComplete, + FinishedIncomplete, +} + +/// A raw iterator over the storage. +pub struct RawIter +where + H: Hasher, +{ + stop_on_incomplete_database: bool, + root: H::Out, + child_info: Option, + trie_iter: TrieDBRawIterator>, + state: IterState, + _phantom: PhantomData<(S, C)>, +} + +impl RawIter +where + H: Hasher, + S: TrieBackendStorage, + H::Out: Codec + Ord, + C: AsLocalTrieCache + Send + Sync, +{ + #[inline] + fn prepare( + &mut self, + backend: &TrieBackendEssence, + callback: impl FnOnce( + &sp_trie::TrieDB>, + &mut TrieDBRawIterator>, + ) -> Option::Out>>>>, + ) -> Option> { + if !matches!(self.state, IterState::Pending) { + return None + } + + let result = backend.with_trie_db(self.root, self.child_info.as_ref(), |db| { + callback(&db, &mut self.trie_iter) + }); + match result { + Some(Ok(key_value)) => Some(Ok(key_value)), + None => { + self.state = IterState::FinishedComplete; + None + }, + Some(Err(error)) => { + self.state = IterState::FinishedIncomplete; + if matches!(*error, TrieError::IncompleteDatabase(_)) && + self.stop_on_incomplete_database + { + None + } else { + Some(Err(format!("TrieDB iteration error: {}", error))) + } + }, + } + } +} + +impl Default for RawIter +where + H: Hasher, +{ + fn default() -> Self { + Self { + stop_on_incomplete_database: false, + child_info: None, + root: Default::default(), + trie_iter: TrieDBRawIterator::empty(), + state: IterState::FinishedComplete, + _phantom: Default::default(), + } + } +} + +impl StorageIterator for RawIter +where + H: Hasher, + S: TrieBackendStorage, + H::Out: Codec + Ord, + C: AsLocalTrieCache + Send + Sync, +{ + type Backend = crate::TrieBackend; + type Error = crate::DefaultError; + + #[inline] + fn next_key(&mut self, backend: &Self::Backend) -> Option> { + self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_key(&trie)) + } + + #[inline] + fn next_pair(&mut self, backend: &Self::Backend) -> Option> { + self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_item(&trie)) + } + + fn was_complete(&self) -> bool { + matches!(self.state, IterState::FinishedComplete) + } +} + /// Patricia trie-based pairs storage essence. pub struct TrieBackendEssence, H: Hasher, C> { storage: S, @@ -258,109 +361,6 @@ impl, H: Hasher, C: AsLocalTrieCache> TrieBackendEss } } -enum IterState { - Pending, - FinishedComplete, - FinishedIncomplete, -} - -/// A raw iterator over the storage. -pub struct RawIter -where - H: Hasher, -{ - stop_on_incomplete_database: bool, - root: H::Out, - child_info: Option, - trie_iter: TrieDBRawIterator>, - state: IterState, - _phantom: PhantomData<(S, C)>, -} - -impl RawIter -where - H: Hasher, - S: TrieBackendStorage, - H::Out: Codec + Ord, - C: AsLocalTrieCache + Send + Sync, -{ - #[inline] - fn prepare( - &mut self, - backend: &TrieBackendEssence, - callback: impl FnOnce( - &sp_trie::TrieDB>, - &mut TrieDBRawIterator>, - ) -> Option::Out>>>>, - ) -> Option> { - if !matches!(self.state, IterState::Pending) { - return None - } - - let result = backend.with_trie_db(self.root, self.child_info.as_ref(), |db| { - callback(&db, &mut self.trie_iter) - }); - match result { - Some(Ok(key_value)) => Some(Ok(key_value)), - None => { - self.state = IterState::FinishedComplete; - None - }, - Some(Err(error)) => { - self.state = IterState::FinishedIncomplete; - if matches!(*error, TrieError::IncompleteDatabase(_)) && - self.stop_on_incomplete_database - { - None - } else { - Some(Err(format!("TrieDB iteration error: {}", error))) - } - }, - } - } -} - -impl Default for RawIter -where - H: Hasher, -{ - fn default() -> Self { - Self { - stop_on_incomplete_database: false, - child_info: None, - root: Default::default(), - trie_iter: TrieDBRawIterator::empty(), - state: IterState::FinishedComplete, - _phantom: Default::default(), - } - } -} - -impl StorageIterator for RawIter -where - H: Hasher, - S: TrieBackendStorage, - H::Out: Codec + Ord, - C: AsLocalTrieCache + Send + Sync, -{ - type Backend = crate::TrieBackend; - type Error = crate::DefaultError; - - #[inline] - fn next_key(&mut self, backend: &Self::Backend) -> Option> { - self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_key(&trie)) - } - - #[inline] - fn next_pair(&mut self, backend: &Self::Backend) -> Option> { - self.prepare(&backend.essence, |trie, trie_iter| trie_iter.next_item(&trie)) - } - - fn was_complete(&self) -> bool { - matches!(self.state, IterState::FinishedComplete) - } -} - impl, H: Hasher, C: AsLocalTrieCache + Send + Sync> TrieBackendEssence where From 2516435116f9b68c3421342b091b3b87bce6659a Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 08:24:24 +0000 Subject: [PATCH 20/22] Remove a TODO in tests --- primitives/state-machine/src/testing.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/primitives/state-machine/src/testing.rs b/primitives/state-machine/src/testing.rs index 0e5e58bbecaae..211a5d1617965 100644 --- a/primitives/state-machine/src/testing.rs +++ b/primitives/state-machine/src/testing.rs @@ -245,7 +245,7 @@ where .backend .pairs(Default::default()) .expect("creating an iterator over all of the pairs doesn't fail in tests") - .collect(); // TODO: Print those out without a `collect`. + .collect(); write!(f, "overlay: {:?}\nbackend: {:?}", self.overlay, pairs) } } From fa585d984040f9eda300366dfb52055fe1bdb3aa Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Mon, 13 Feb 2023 08:34:42 +0000 Subject: [PATCH 21/22] Update comment for `StorageIterator::was_complete` --- primitives/state-machine/src/backend.rs | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/primitives/state-machine/src/backend.rs b/primitives/state-machine/src/backend.rs index 1c098eb6e8602..bf4372fbd2a68 100644 --- a/primitives/state-machine/src/backend.rs +++ b/primitives/state-machine/src/backend.rs @@ -77,9 +77,7 @@ where backend: &Self::Backend, ) -> Option>; - /// Returns whether the trie end was reached. - /// - /// Will return `true` once we've reached the end of the trie. + /// Returns whether the end of iteration was reached without an error. fn was_complete(&self) -> bool; } From fdb6b635e3c2f7af794e9522123109b4e4c68cb9 Mon Sep 17 00:00:00 2001 From: Jan Bujak Date: Tue, 14 Feb 2023 14:39:07 +0000 Subject: [PATCH 22/22] Update `trie-db` to 0.25.1 --- Cargo.lock | 4 ++-- primitives/state-machine/Cargo.toml | 2 +- test-utils/runtime/Cargo.toml | 2 +- utils/frame/rpc/state-trie-migration-rpc/Cargo.toml | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8db9e8a8c617b..7bcb5211fa55e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11345,9 +11345,9 @@ dependencies = [ [[package]] name = "trie-db" -version = "0.25.0" +version = "0.25.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9bcd4e7fbb3bcab17b02596b53876e36eb39663ddd87884bfd88c69cdeb2ebd6" +checksum = "3390c0409daaa6027d6681393316f4ccd3ff82e1590a1e4725014e3ae2bf1920" dependencies = [ "hash-db", "hashbrown 0.13.2", diff --git a/primitives/state-machine/Cargo.toml b/primitives/state-machine/Cargo.toml index dd47c8e988478..2969a51c4eba0 100644 --- a/primitives/state-machine/Cargo.toml +++ b/primitives/state-machine/Cargo.toml @@ -33,7 +33,7 @@ array-bytes = "4.1" pretty_assertions = "1.2.1" rand = "0.8.5" sp-runtime = { version = "7.0.0", path = "../runtime" } -trie-db = "0.25.0" +trie-db = "0.25.1" assert_matches = "1.5" [features] diff --git a/test-utils/runtime/Cargo.toml b/test-utils/runtime/Cargo.toml index 1f96bcbb7bdcc..33dbfce680d75 100644 --- a/test-utils/runtime/Cargo.toml +++ b/test-utils/runtime/Cargo.toml @@ -41,7 +41,7 @@ pallet-timestamp = { version = "4.0.0-dev", default-features = false, path = ".. sp-finality-grandpa = { version = "4.0.0-dev", default-features = false, path = "../../primitives/finality-grandpa" } sp-trie = { version = "7.0.0", default-features = false, path = "../../primitives/trie" } sp-transaction-pool = { version = "4.0.0-dev", default-features = false, path = "../../primitives/transaction-pool" } -trie-db = { version = "0.25.0", default-features = false } +trie-db = { version = "0.25.1", default-features = false } sc-service = { version = "0.10.0-dev", default-features = false, optional = true, features = ["test-helpers"], path = "../../client/service" } sp-state-machine = { version = "0.13.0", default-features = false, path = "../../primitives/state-machine" } sp-externalities = { version = "0.13.0", default-features = false, path = "../../primitives/externalities" } diff --git a/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml b/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml index c7c2f442b56f7..1dd9da9a56fee 100644 --- a/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml +++ b/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml @@ -21,7 +21,7 @@ log = { version = "0.4.17", default-features = false } sp-core = { path = "../../../../primitives/core" } sp-state-machine = { path = "../../../../primitives/state-machine" } sp-trie = { path = "../../../../primitives/trie" } -trie-db = "0.25.0" +trie-db = "0.25.1" jsonrpsee = { version = "0.16.2", features = ["client-core", "server", "macros"] }