Paged keys rpc for child storage. (#9100)

* childstate_getKeysPaged rpc

* Rename `v` to `iter`.

* Update client/api/src/backend.rs

Co-authored-by: Alexander Popiak <[email protected]>

* Update client/api/src/backend.rs

Co-authored-by: Alexander Popiak <[email protected]>

Co-authored-by: Alexander Popiak <[email protected]>

client/api/src/backend.rs

@@ -280,6 +280,7 @@ pub trait AuxStore {
 /// An `Iterator` that iterates keys in a given block under a prefix.
 pub struct KeyIterator<'a, State, Block> {
 	state: State,
+	child_storage: Option<ChildInfo>,
 	prefix: Option<&'a StorageKey>,
 	current_key: Vec<u8>,
 	_phantom: PhantomData<Block>,
@@ -290,6 +291,23 @@ impl <'a, State, Block> KeyIterator<'a, State, Block> {
 	pub fn new(state: State, prefix: Option<&'a StorageKey>, current_key: Vec<u8>) -> Self {
 		Self {
 			state,
+			child_storage: None,
+			prefix,
+			current_key,
+			_phantom: PhantomData,
+		}
+	}
+
+	/// Create a `KeyIterator` instance for a child storage.
+	pub fn new_child(
+		state: State,
+		child_info: ChildInfo,
+		prefix: Option<&'a StorageKey>,
+		current_key: Vec<u8>,
+	) -> Self {
+		Self {
+			state,
+			child_storage: Some(child_info),
 			prefix,
 			current_key,
 			_phantom: PhantomData,
@@ -304,10 +322,11 @@ impl<'a, State, Block> Iterator for KeyIterator<'a, State, Block> where
 	type Item = StorageKey;
 
 	fn next(&mut self) -> Option<Self::Item> {
-		let next_key = self.state
-			.next_storage_key(&self.current_key)
-			.ok()
-			.flatten()?;
+		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(prefix) = self.prefix {
 			if !next_key.starts_with(&prefix.0[..]) {
@@ -361,6 +380,16 @@ pub trait StorageProvider<Block: BlockT, B: Backend<Block>> {
 		key_prefix: &StorageKey
 	) -> sp_blockchain::Result<Vec<StorageKey>>;
 
+	/// Given a `BlockId` 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<'a>(
+		&self,
+		id: &BlockId<Block>,
+		child_info: ChildInfo,
+		prefix: Option<&'a StorageKey>,
+		start_key: Option<&StorageKey>
+	) -> sp_blockchain::Result<KeyIterator<'a, B::State, Block>>;
+
 	/// Given a `BlockId`, a key and a child storage key, return the hash under the key in that block.
 	fn child_storage_hash(
 		&self,

client/rpc-api/src/child_state/mod.rs

@@ -34,6 +34,7 @@ pub trait ChildStateApi<Hash> {
 	/// RPC Metadata
 	type Metadata;
 
+	/// DEPRECATED: Please use `childstate_getKeysPaged` with proper paging support.
 	/// Returns the keys with prefix from a child storage, leave empty to get all the keys
 	#[rpc(name = "childstate_getKeys")]
 	fn storage_keys(
@@ -43,6 +44,19 @@ pub trait ChildStateApi<Hash> {
 		hash: Option<Hash>
 	) -> FutureResult<Vec<StorageKey>>;
 
+	/// Returns the keys with prefix from a child storage with pagination support.
+	/// Up to `count` keys will be returned.
+	/// If `start_key` is passed, return next keys in storage in lexicographic order.
+	#[rpc(name = "childstate_getKeysPaged", alias("childstate_getKeysPagedAt"))]
+	fn storage_keys_paged(
+		&self,
+		child_storage_key: PrefixedStorageKey,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+		hash: Option<Hash>,
+	) -> FutureResult<Vec<StorageKey>>;
+
 	/// Returns a child storage entry at a specific block's state.
 	#[rpc(name = "childstate_getStorage")]
 	fn storage(

client/rpc/src/state/mod.rs

@@ -402,6 +402,16 @@ pub trait ChildStateBackend<Block: BlockT, Client>: Send + Sync + 'static
 		prefix: StorageKey,
 	) -> FutureResult<Vec<StorageKey>>;
 
+	/// Returns the keys with prefix from a child storage with pagination support.
+	fn storage_keys_paged(
+		&self,
+		block: Option<Block::Hash>,
+		storage_key: PrefixedStorageKey,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+	) -> FutureResult<Vec<StorageKey>>;
+
 	/// Returns a child storage entry at a specific block's state.
 	fn storage(
 		&self,
@@ -469,6 +479,17 @@ impl<Block, Client> ChildStateApi<Block::Hash> for ChildState<Block, Client>
 		self.backend.storage_keys(block, storage_key, key_prefix)
 	}
 
+	fn storage_keys_paged(
+		&self,
+		storage_key: PrefixedStorageKey,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+		block: Option<Block::Hash>,
+	) -> FutureResult<Vec<StorageKey>> {
+		self.backend.storage_keys_paged(block, storage_key, prefix, count, start_key)
+	}
+
 	fn storage_hash(
 		&self,
 		storage_key: PrefixedStorageKey,

client/rpc/src/state/state_full.rs

@@ -296,7 +296,7 @@ impl<BE, Block, Client> StateBackend<Block, Client> for FullState<BE, Block, Cli
 						&BlockId::Hash(block), prefix.as_ref(), start_key.as_ref()
 					)
 				)
-				.map(|v| v.take(count as usize).collect())
+				.map(|iter| iter.take(count as usize).collect())
 				.map_err(client_err)))
 	}
 
@@ -618,6 +618,29 @@ impl<BE, Block, Client> ChildStateBackend<Block, Client> for FullState<BE, Block
 				.map_err(client_err)))
 	}
 
+	fn storage_keys_paged(
+		&self,
+		block: Option<Block::Hash>,
+		storage_key: PrefixedStorageKey,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+	) -> FutureResult<Vec<StorageKey>> {
+		Box::new(result(
+			self.block_or_best(block)
+				.and_then(|block| {
+					let child_info = match ChildType::from_prefixed_key(&storage_key) {
+						Some((ChildType::ParentKeyId, storage_key)) => ChildInfo::new_default(storage_key),
+						None => return Err(sp_blockchain::Error::InvalidChildStorageKey),
+					};
+					self.client.child_storage_keys_iter(
+						&BlockId::Hash(block), child_info, prefix.as_ref(), start_key.as_ref(),
+					)
+				})
+				.map(|iter| iter.take(count as usize).collect())
+				.map_err(client_err)))
+	}
+
 	fn storage(
 		&self,
 		block: Option<Block::Hash>,

client/rpc/src/state/state_light.rs

@@ -509,6 +509,17 @@ impl<Block, F, Client> ChildStateBackend<Block, Client> for LightState<Block, F,
 		Box::new(result(Err(client_err(ClientError::NotAvailableOnLightClient))))
 	}
 
+	fn storage_keys_paged(
+		&self,
+		_block: Option<Block::Hash>,
+		_storage_key: PrefixedStorageKey,
+		_prefix: Option<StorageKey>,
+		_count: u32,
+		_start_key: Option<StorageKey>,
+	) -> FutureResult<Vec<StorageKey>> {
+		Box::new(result(Err(client_err(ClientError::NotAvailableOnLightClient))))
+	}
+
 	fn storage(
 		&self,
 		block: Option<Block::Hash>,

client/service/src/client/client.rs

@@ -1481,7 +1481,6 @@ impl<B, E, Block, RA> StorageProvider<Block, B> for Client<B, E, Block, RA> wher
 		Ok(keys)
 	}
 
-
 	fn storage_keys_iter<'a>(
 		&self,
 		id: &BlockId<Block>,
@@ -1496,6 +1495,20 @@ impl<B, E, Block, RA> StorageProvider<Block, B> for Client<B, E, Block, RA> wher
 		Ok(KeyIterator::new(state, prefix, start_key))
 	}
 
+	fn child_storage_keys_iter<'a>(
+		&self,
+		id: &BlockId<Block>,
+		child_info: ChildInfo,
+		prefix: Option<&'a StorageKey>,
+		start_key: Option<&StorageKey>
+	) -> sp_blockchain::Result<KeyIterator<'a, B::State, Block>> {
+		let state = self.state_at(id)?;
+		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, start_key))
+	}
 
 	fn storage(
 		&self,

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

@@ -50,7 +50,7 @@ use sp_consensus::{
 	BlockOrigin, SelectChain, BlockImport, Error as ConsensusError, BlockCheckParams, ImportResult,
 	BlockStatus, BlockImportParams, ForkChoiceStrategy,
 };
-use sp_storage::StorageKey;
+use sp_storage::{StorageKey, ChildInfo};
 use sp_trie::{TrieConfiguration, trie_types::Layout};
 use sp_runtime::{generic::BlockId, DigestItem, Justifications};
 use hex_literal::hex;
@@ -1999,15 +1999,26 @@ fn imports_blocks_with_changes_tries_config_change() {
 
 #[test]
 fn storage_keys_iter_prefix_and_start_key_works() {
-	let client = substrate_test_runtime_client::new();
-
+	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])
+		.add_extra_child_storage(&child_info, b"second".to_vec(), vec![0u8; 32])
+		.add_extra_child_storage(&child_info, b"third".to_vec(), vec![0u8; 32])
+		.build();
+
+	let child_root = b":child_storage:default:child".to_vec();
 	let prefix = StorageKey(hex!("3a").to_vec());
+	let child_prefix = StorageKey(b"sec".to_vec());
 
 	let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), None)
 		.unwrap()
 		.map(|x| x.0)
 		.collect();
-	assert_eq!(res, [hex!("3a636f6465").to_vec(), hex!("3a686561707061676573").to_vec()]);
+	assert_eq!(res, [
+		child_root.clone(),
+		hex!("3a636f6465").to_vec(),
+		hex!("3a686561707061676573").to_vec(),
+	]);
 
 	let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), Some(&StorageKey(hex!("3a636f6465").to_vec())))
 		.unwrap()
@@ -2020,6 +2031,26 @@ fn storage_keys_iter_prefix_and_start_key_works() {
 		.map(|x| x.0)
 		.collect();
 	assert_eq!(res, Vec::<Vec<u8>>::new());
+
+	let res: Vec<_> = client.child_storage_keys_iter(
+		&BlockId::Number(0),
+		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(
+		&BlockId::Number(0),
+		child_info,
+		None,
+		Some(&StorageKey(b"second".to_vec())),
+	).unwrap()
+		.map(|x| x.0)
+		.collect();
+	assert_eq!(res, [b"third".to_vec()]);
 }
 
 #[test]