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`

substrate/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(

substrate/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,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())));