Speed up storage iteration from within the runtime (paritytech#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`

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::<T>::insert(i, sp_std::vec![0; 64]);
+		}
+	}: {
+		for (key, value) in UnboundedMapTwox::<T>::iter() {
+			unsafe {
+				core::ptr::read_volatile(&key);
+				core::ptr::read_volatile(value.as_ptr());
+			}
+		}
+	}
+
 	impl_benchmark_test_suite!(
 		Pallet,
 		mock::new_test_ext(),

frame/benchmarking/pov/src/lib.rs

@@ -107,6 +107,11 @@ pub mod pallet {
 	pub(crate) type UnboundedMap2<T: Config> =
 		StorageMap<Hasher = Blake2_256, Key = u32, Value = Vec<u32>, QueryKind = OptionQuery>;
 
+	#[pallet::storage]
+	#[pallet::unbounded]
+	pub(crate) type UnboundedMapTwox<T: Config> =
+		StorageMap<Hasher = Twox64Concat, Key = u32, Value = Vec<u32>, QueryKind = OptionQuery>;
+
 	#[pallet::event]
 	#[pallet::generate_deposit(pub(super) fn deposit_event)]
 	pub enum Event<T: Config> {

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]

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<S, H, C> {
 		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<S, H, C>
+where
+	H: Hasher,
+{
+	last_key: sp_std::vec::Vec<u8>,
+	iter: RawIter<S, H, C>,
+}
+
+impl<S, H, C> Default for CachedIter<S, H, C>
+where
+	H: Hasher,
+{
+	fn default() -> Self {
+		Self { last_key: Default::default(), iter: Default::default() }
 	}
 }
 
+#[cfg(feature = "std")]
+type CacheCell<T> = parking_lot::Mutex<T>;
+
+#[cfg(not(feature = "std"))]
+type CacheCell<T> = core::cell::RefCell<T>;
+
+#[cfg(feature = "std")]
+fn access_cache<T, R>(cell: &CacheCell<T>, callback: impl FnOnce(&mut T) -> R) -> R {
+	callback(&mut *cell.lock())
+}
+
+#[cfg(not(feature = "std"))]
+fn access_cache<T, R>(cell: &CacheCell<T>, 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<S: TrieBackendStorage<H>, H: Hasher, C = LocalTrieCache<H>> {
 	pub(crate) essence: TrieBackendEssence<S, H, C>,
+	next_storage_key_cache: CacheCell<Option<CachedIter<S, H, C>>>,
 }
 
 impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + Send + Sync> TrieBackend<S, H, C>
 where
 	H::Out: Codec,
 {
+	#[cfg(test)]
+	pub(crate) fn from_essence(essence: TrieBackendEssence<S, H, C>) -> Self {
+		Self { essence, next_storage_key_cache: Default::default() }
+	}
+
 	/// Get backend essence reference.
 	pub fn essence(&self) -> &TrieBackendEssence<S, H, C> {
 		&self.essence
@@ -265,7 +309,40 @@ where
 	}
 
 	fn next_storage_key(&self, key: &[u8]) -> Result<Option<StorageKey>, 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(

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<Option<StorageKey>> {
+	///
+	/// 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<Option<StorageKey>> {
 		self.next_storage_key_from_root(&self.root, None, key)
 	}
 
@@ -859,6 +864,7 @@ impl<S: TrieBackendStorage<H>, H: Hasher, C: AsLocalTrieCache<H> + 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,14 +903,15 @@ 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())));
 		assert_eq!(essence_1.next_storage_key(b"4"), Ok(Some(b"6".to_vec())));
 		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())));

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" }
 

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" }

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"] }