From 9c6c1d52bdf726642be166a68483a710aeddd321 Mon Sep 17 00:00:00 2001 From: Koute Date: Wed, 1 Mar 2023 17:58:18 +0900 Subject: [PATCH] Speed up storage iteration from within the runtime (#13479) * Speed up storage iteration from within the runtime * Move the cached iterator into an `Option` * Use `RefCell` in no_std * Simplify the code slightly * Use `Option::replace` * Update doc comment for `next_storage_key_slow` --- Cargo.lock | 8 +- frame/benchmarking/pov/src/benchmarking.rs | 13 +++ frame/benchmarking/pov/src/lib.rs | 5 ++ primitives/state-machine/Cargo.toml | 2 +- primitives/state-machine/src/trie_backend.rs | 85 ++++++++++++++++++- .../state-machine/src/trie_backend_essence.rs | 13 ++- primitives/trie/Cargo.toml | 4 +- test-utils/runtime/Cargo.toml | 2 +- .../rpc/state-trie-migration-rpc/Cargo.toml | 2 +- 9 files changed, 118 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a412347bde5c0..a8f596c0684ef 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11437,9 +11437,9 @@ dependencies = [ [[package]] name = "trie-bench" -version = "0.35.0" +version = "0.36.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22c1d18c423077531e693e87ace54ed9e4af1e4ce0a3ea8c9aa6608741074e2b" +checksum = "ac2b7695feb8041efc0adaa09ed3a692ca7b7c997395954fdc838b5b078346f6" dependencies = [ "criterion", "hash-db", @@ -11453,9 +11453,9 @@ dependencies = [ [[package]] name = "trie-db" -version = "0.25.1" +version = "0.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3390c0409daaa6027d6681393316f4ccd3ff82e1590a1e4725014e3ae2bf1920" +checksum = "879380c0061b165ba1f036325b7f3478bc1a947814d9fc36d22c5d8e960b11bd" dependencies = [ "hash-db", "hashbrown 0.13.2", diff --git a/frame/benchmarking/pov/src/benchmarking.rs b/frame/benchmarking/pov/src/benchmarking.rs index ebe9b241416ad..46a98ae3ff183 100644 --- a/frame/benchmarking/pov/src/benchmarking.rs +++ b/frame/benchmarking/pov/src/benchmarking.rs @@ -316,6 +316,19 @@ frame_benchmarking::benchmarks! { call.dispatch_bypass_filter(RawOrigin::Root.into()).unwrap(); } + storage_iteration { + for i in 0..65000 { + UnboundedMapTwox::::insert(i, sp_std::vec![0; 64]); + } + }: { + for (key, value) in UnboundedMapTwox::::iter() { + unsafe { + core::ptr::read_volatile(&key); + core::ptr::read_volatile(value.as_ptr()); + } + } + } + impl_benchmark_test_suite!( Pallet, mock::new_test_ext(), diff --git a/frame/benchmarking/pov/src/lib.rs b/frame/benchmarking/pov/src/lib.rs index 7dee5c37c1288..b66b5e417817a 100644 --- a/frame/benchmarking/pov/src/lib.rs +++ b/frame/benchmarking/pov/src/lib.rs @@ -107,6 +107,11 @@ pub mod pallet { pub(crate) type UnboundedMap2 = StorageMap, QueryKind = OptionQuery>; + #[pallet::storage] + #[pallet::unbounded] + pub(crate) type UnboundedMapTwox = + StorageMap, QueryKind = OptionQuery>; + #[pallet::event] #[pallet::generate_deposit(pub(super) fn deposit_event)] pub enum Event { diff --git a/primitives/state-machine/Cargo.toml b/primitives/state-machine/Cargo.toml index 2969a51c4eba0..56fbe0726b82e 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.1" +trie-db = "0.26.0" assert_matches = "1.5" [features] diff --git a/primitives/state-machine/src/trie_backend.rs b/primitives/state-machine/src/trie_backend.rs index 3e8d0a7a3bf0a..10e2dfd16d66a 100644 --- a/primitives/state-machine/src/trie_backend.rs +++ b/primitives/state-machine/src/trie_backend.rs @@ -20,8 +20,8 @@ #[cfg(feature = "std")] use crate::backend::AsTrieBackend; use crate::{ - backend::IterArgs, - trie_backend_essence::{TrieBackendEssence, TrieBackendStorage}, + backend::{IterArgs, StorageIterator}, + trie_backend_essence::{RawIter, TrieBackendEssence, TrieBackendStorage}, Backend, StorageKey, StorageValue, }; use codec::Codec; @@ -172,6 +172,7 @@ where self.cache, self.recorder, ), + next_storage_key_cache: Default::default(), } } @@ -180,19 +181,62 @@ where pub fn build(self) -> TrieBackend { let _ = self.cache; - TrieBackend { essence: TrieBackendEssence::new(self.storage, self.root) } + TrieBackend { + essence: TrieBackendEssence::new(self.storage, self.root), + next_storage_key_cache: Default::default(), + } + } +} + +/// A cached iterator. +struct CachedIter +where + H: Hasher, +{ + last_key: sp_std::vec::Vec, + iter: RawIter, +} + +impl Default for CachedIter +where + H: Hasher, +{ + fn default() -> Self { + Self { last_key: Default::default(), iter: Default::default() } } } +#[cfg(feature = "std")] +type CacheCell = parking_lot::Mutex; + +#[cfg(not(feature = "std"))] +type CacheCell = core::cell::RefCell; + +#[cfg(feature = "std")] +fn access_cache(cell: &CacheCell, callback: impl FnOnce(&mut T) -> R) -> R { + callback(&mut *cell.lock()) +} + +#[cfg(not(feature = "std"))] +fn access_cache(cell: &CacheCell, callback: impl FnOnce(&mut T) -> R) -> R { + callback(&mut *cell.borrow_mut()) +} + /// Patricia trie-based backend. Transaction type is an overlay of changes to commit. pub struct TrieBackend, H: Hasher, C = LocalTrieCache> { pub(crate) essence: TrieBackendEssence, + next_storage_key_cache: CacheCell>>, } impl, H: Hasher, C: AsLocalTrieCache + Send + Sync> TrieBackend where H::Out: Codec, { + #[cfg(test)] + pub(crate) fn from_essence(essence: TrieBackendEssence) -> Self { + Self { essence, next_storage_key_cache: Default::default() } + } + /// Get backend essence reference. pub fn essence(&self) -> &TrieBackendEssence { &self.essence @@ -265,7 +309,40 @@ where } fn next_storage_key(&self, key: &[u8]) -> Result, Self::Error> { - self.essence.next_storage_key(key) + let (is_cached, mut cache) = access_cache(&self.next_storage_key_cache, Option::take) + .map(|cache| (cache.last_key == key, cache)) + .unwrap_or_default(); + + if !is_cached { + cache.iter = self.raw_iter(IterArgs { + start_at: Some(key), + start_at_exclusive: true, + ..IterArgs::default() + })? + }; + + let next_key = match cache.iter.next_key(self) { + None => return Ok(None), + Some(Err(error)) => return Err(error), + Some(Ok(next_key)) => next_key, + }; + + cache.last_key.clear(); + cache.last_key.extend_from_slice(&next_key); + access_cache(&self.next_storage_key_cache, |cache_cell| cache_cell.replace(cache)); + + #[cfg(debug_assertions)] + debug_assert_eq!( + self.essence + .next_storage_key_slow(key) + .expect( + "fetching the next key through iterator didn't fail so this shouldn't either" + ) + .as_ref(), + Some(&next_key) + ); + + Ok(Some(next_key)) } fn next_child_storage_key( diff --git a/primitives/state-machine/src/trie_backend_essence.rs b/primitives/state-machine/src/trie_backend_essence.rs index f071a32cede89..1f6d71b2dce80 100644 --- a/primitives/state-machine/src/trie_backend_essence.rs +++ b/primitives/state-machine/src/trie_backend_essence.rs @@ -415,9 +415,14 @@ where }) } - /// Return the next key in the trie i.e. the minimum key that is strictly superior to `key` in + /// Returns 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> { + /// + /// Will always traverse the trie from scratch in search of the key, which is slow. + /// Used only when debug assertions are enabled to crosscheck the results of finding + /// the next key through an iterator. + #[cfg(debug_assertions)] + pub fn next_storage_key_slow(&self, key: &[u8]) -> Result> { self.next_storage_key_from_root(&self.root, None, key) } @@ -859,6 +864,7 @@ impl, H: Hasher, C: AsLocalTrieCache + Send + Sync> #[cfg(test)] mod test { use super::*; + use crate::{Backend, TrieBackend}; use sp_core::{Blake2Hasher, H256}; use sp_trie::{ cache::LocalTrieCache, trie_types::TrieDBMutBuilderV1 as TrieDBMutBuilder, KeySpacedDBMut, @@ -897,6 +903,8 @@ mod test { }; let essence_1 = TrieBackendEssence::<_, _, LocalTrieCache<_>>::new(mdb, root_1); + let mdb = essence_1.backend_storage().clone(); + let essence_1 = TrieBackend::from_essence(essence_1); assert_eq!(essence_1.next_storage_key(b"2"), Ok(Some(b"3".to_vec()))); assert_eq!(essence_1.next_storage_key(b"3"), Ok(Some(b"4".to_vec()))); @@ -904,7 +912,6 @@ mod test { assert_eq!(essence_1.next_storage_key(b"5"), Ok(Some(b"6".to_vec()))); assert_eq!(essence_1.next_storage_key(b"6"), Ok(None)); - let mdb = essence_1.backend_storage().clone(); let essence_2 = TrieBackendEssence::<_, _, LocalTrieCache<_>>::new(mdb, root_2); assert_eq!(essence_2.next_child_storage_key(child_info, b"2"), Ok(Some(b"3".to_vec()))); diff --git a/primitives/trie/Cargo.toml b/primitives/trie/Cargo.toml index 33a62cdd94084..fae39ec34c0ea 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.25.0", default-features = false } +trie-db = { version = "0.26.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" } @@ -38,7 +38,7 @@ schnellru = { version = "0.2.1", optional = true } [dev-dependencies] array-bytes = "4.1" criterion = "0.4.0" -trie-bench = "0.35.0" +trie-bench = "0.36.0" trie-standardmap = "0.15.2" sp-runtime = { version = "7.0.0", path = "../runtime" } diff --git a/test-utils/runtime/Cargo.toml b/test-utils/runtime/Cargo.toml index 339a435ef2a04..e0f414b161cf5 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-consensus-grandpa = { version = "4.0.0-dev", default-features = false, path = "../../primitives/consensus/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.1", default-features = false } +trie-db = { version = "0.26.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/rpc/state-trie-migration-rpc/Cargo.toml b/utils/frame/rpc/state-trie-migration-rpc/Cargo.toml index 1dd9da9a56fee..e6eccc6ff62e0 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.1" +trie-db = "0.26.0" jsonrpsee = { version = "0.16.2", features = ["client-core", "server", "macros"] }