Rework storage iterators (#13284)

* Rework storage iterators

* Make sure storage iteration is also accounted for when benchmarking

* Use `trie-db` from crates.io

* Appease clippy

* Bump `trie-bench` to 0.35.0

* Fix tests' compilation

* Update comment to clarify how `IterArgs::start_at` works

* Add extra tests

* Fix iterators on `Client` so that they behave as before

* Add extra `unwrap`s in tests

* More clippy fixes

* Come on clippy, give me a break already

* Rename `allow_missing` to `stop_on_incomplete_database`

* Add `#[inline]` to `with_recorder_and_cache`

* Use `with_recorder_and_cache` in `with_trie_db`; add doc comment

* Simplify code: use `with_trie_db` in `next_storage_key_from_root`

* Remove `expect`s in the benchmarking CLI

* Add extra doc comments

* Move `RawIter` before `TrieBackendEssence` (no code changes; just cut-paste)

* Remove a TODO in tests

* Update comment for `StorageIterator::was_complete`

* Update `trie-db` to 0.25.1

substrate/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<B>,
 {
-	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 })
 }

substrate/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,
@@ -1462,52 +1462,37 @@ where
 	Block: BlockT,
 {
 	fn storage_keys(
-		&self,
-		hash: Block::Hash,
-		key_prefix: &StorageKey,
-	) -> sp_blockchain::Result<Vec<StorageKey>> {
-		let keys = self.state_at(hash)?.keys(&key_prefix.0).into_iter().map(StorageKey).collect();
-		Ok(keys)
-	}
-
-	fn storage_pairs(
-		&self,
-		hash: <Block as BlockT>::Hash,
-		key_prefix: &StorageKey,
-	) -> sp_blockchain::Result<Vec<(StorageKey, StorageData)>> {
-		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)
-	}
-
-	fn storage_keys_iter(
 		&self,
 		hash: <Block as BlockT>::Hash,
 		prefix: Option<&StorageKey>,
 		start_key: Option<&StorageKey>,
-	) -> sp_blockchain::Result<KeyIterator<B::State, Block>> {
+	) -> sp_blockchain::Result<KeysIter<B::State, Block>> {
 		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))
+		KeysIter::new(state, prefix, start_key)
+			.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))
 	}
 
-	fn child_storage_keys_iter(
+	fn child_storage_keys(
 		&self,
 		hash: <Block as BlockT>::Hash,
 		child_info: ChildInfo,
 		prefix: Option<&StorageKey>,
 		start_key: Option<&StorageKey>,
-	) -> sp_blockchain::Result<KeyIterator<B::State, Block>> {
+	) -> sp_blockchain::Result<KeysIter<B::State, Block>> {
 		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))
+		KeysIter::new_child(state, child_info, prefix, start_key)
+			.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))
+	}
+
+	fn storage_pairs(
+		&self,
+		hash: <Block as BlockT>::Hash,
+		prefix: Option<&StorageKey>,
+		start_key: Option<&StorageKey>,
+	) -> sp_blockchain::Result<PairsIter<B::State, Block>> {
+		let state = self.state_at(hash)?;
+		PairsIter::new(state, prefix, start_key)
+			.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))
 	}
 
 	fn storage(
@@ -1532,21 +1517,6 @@ where
 			.map_err(|e| sp_blockchain::Error::from_state(Box::new(e)))
 	}
 
-	fn child_storage_keys(
-		&self,
-		hash: <Block as BlockT>::Hash,
-		child_info: &ChildInfo,
-		key_prefix: &StorageKey,
-	) -> sp_blockchain::Result<Vec<StorageKey>> {
-		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: <Block as BlockT>::Hash,

substrate/client/service/test/src/client/mod.rs

@@ -341,7 +341,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::<Vec<_>>(),
+		client
+			.state_at(hash0)
+			.unwrap()
+			.pairs(Default::default())
+			.unwrap()
+			.collect::<Vec<_>>()
+	);
 	assert_eq!(
 		client
 			.runtime_api()
@@ -394,8 +407,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::<Vec<_>>(),
+		client
+			.state_at(hashof0)
+			.unwrap()
+			.pairs(Default::default())
+			.unwrap()
+			.collect::<Vec<_>>()
 	);
 	assert_eq!(client.body(hashof1).unwrap().unwrap().len(), 1)
 }
@@ -1688,7 +1711,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])
@@ -1703,7 +1726,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();
@@ -1717,7 +1740,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"))),
@@ -1728,7 +1751,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"))),
@@ -1739,19 +1762,14 @@ fn storage_keys_iter_prefix_and_start_key_works() {
 	assert_eq!(res, Vec::<Vec<u8>>::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();
@@ -1759,7 +1777,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;
@@ -1767,7 +1785,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))
@@ -1787,8 +1805,56 @@ fn storage_keys_iter_works() {
 		]
 	);
 
+	// Starting at an empty key nothing gets skipped.
 	let res: Vec<_> = client
-		.storage_keys_iter(
+		.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,
 			Some(&prefix),
 			Some(&StorageKey(array_bytes::hex2bytes_unchecked("3a636f6465"))),
@@ -1811,7 +1877,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(