add state_getKeysPaged (#4718)

* add storage_getNextKey

* RPC storage_getKeysPages

* respect count

* use iterator

* improve

* add tests

* improve

* add doc comments

* Make prefix optional

* update error

* improve

substrate/client/rpc-api/src/state/error.rs

@@ -41,6 +41,14 @@ pub enum Error {
 		/// Details of the error message.
 		details: String,
 	},
+	/// Provided count exceeds maximum value.
+	#[display(fmt = "count exceeds maximum value. value: {}, max: {}", value, max)]
+	InvalidCount {
+		/// Provided value
+		value: u32,
+		/// Maximum allowed value
+		max: u32,
+	},
 }
 
 impl std::error::Error for Error {
@@ -63,6 +71,11 @@ impl From<Error> for rpc::Error {
 				message: format!("{}", e),
 				data: None,
 			},
+			Error::InvalidCount { .. } => rpc::Error {
+				code: rpc::ErrorCode::ServerError(BASE_ERROR + 2),
+				message: format!("{}", e),
+				data: None,
+			},
 			e => errors::internal(e),
 		}
 	}

substrate/client/rpc-api/src/state/mod.rs

@@ -39,10 +39,23 @@ pub trait StateApi<Hash> {
 	#[rpc(name = "state_call", alias("state_callAt"))]
 	fn call(&self, name: String, bytes: Bytes, hash: Option<Hash>) -> FutureResult<Bytes>;
 
-	/// Returns the keys with prefix, leave empty to get all the keys
+	/// DEPRECATED: Please use `state_getKeysPaged` with proper paging support.
+	/// Returns the keys with prefix, leave empty to get all the keys.
 	#[rpc(name = "state_getKeys")]
 	fn storage_keys(&self, prefix: StorageKey, hash: Option<Hash>) -> FutureResult<Vec<StorageKey>>;
 
+	/// Returns the keys with prefix 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 = "state_getKeysPaged", alias("state_getKeysPagedAt"))]
+	fn storage_keys_paged(
+		&self,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+		hash: Option<Hash>,
+	) -> FutureResult<Vec<StorageKey>>;
+
 	/// Returns a storage entry at a specific block's state.
 	#[rpc(name = "state_getStorage", alias("state_getStorageAt"))]
 	fn storage(&self, key: StorageKey, hash: Option<Hash>) -> FutureResult<Option<StorageData>>;

substrate/client/rpc/src/state/mod.rs

@@ -24,7 +24,7 @@ mod tests;
 
 use std::sync::Arc;
 use jsonrpc_pubsub::{typed::Subscriber, SubscriptionId};
-use rpc::{Result as RpcResult, futures::Future};
+use rpc::{Result as RpcResult, futures::{Future, future::result}};
 
 use sc_rpc_api::Subscriptions;
 use sc_client::{Client, CallExecutor, light::{blockchain::RemoteBlockchain, fetcher::Fetcher}};
@@ -38,6 +38,8 @@ use self::error::{Error, FutureResult};
 
 pub use sc_rpc_api::state::*;
 
+const STORAGE_KEYS_PAGED_MAX_COUNT: u32 = 1000;
+
 /// State backend API.
 pub trait StateBackend<B, E, Block: BlockT, RA>: Send + Sync + 'static
 	where
@@ -61,6 +63,15 @@ pub trait StateBackend<B, E, Block: BlockT, RA>: Send + Sync + 'static
 		prefix: StorageKey,
 	) -> FutureResult<Vec<StorageKey>>;
 
+	/// Returns the keys with prefix with pagination support.
+	fn storage_keys_paged(
+		&self,
+		block: Option<Block::Hash>,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+	) -> FutureResult<Vec<StorageKey>>;
+
 	/// Returns a storage entry at a specific block's state.
 	fn storage(
 		&self,
@@ -244,6 +255,24 @@ impl<B, E, Block, RA> StateApi<Block::Hash> for State<B, E, Block, RA>
 		self.backend.storage_keys(block, key_prefix)
 	}
 
+	fn storage_keys_paged(
+		&self,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+		block: Option<Block::Hash>,
+	) -> FutureResult<Vec<StorageKey>> {
+		if count > STORAGE_KEYS_PAGED_MAX_COUNT {
+			return Box::new(result(Err(
+				Error::InvalidCount {
+					value: count,
+					max: STORAGE_KEYS_PAGED_MAX_COUNT,
+				}
+			)));
+		}
+		self.backend.storage_keys_paged(block, prefix, count, start_key)
+	}
+
 	fn storage(&self, key: StorageKey, block: Option<Block::Hash>) -> FutureResult<Option<StorageData>> {
 		self.backend.storage(block, key)
 	}

substrate/client/rpc/src/state/state_full.rs

@@ -254,6 +254,24 @@ impl<B, E, Block, RA> StateBackend<B, E, Block, RA> for FullState<B, E, Block, R
 				.map_err(client_err)))
 	}
 
+	fn storage_keys_paged(
+		&self,
+		block: Option<Block::Hash>,
+		prefix: Option<StorageKey>,
+		count: u32,
+		start_key: Option<StorageKey>,
+	) -> FutureResult<Vec<StorageKey>> {
+		Box::new(result(
+			self.block_or_best(block)
+				.and_then(|block|
+					self.client.storage_keys_iter(
+						&BlockId::Hash(block), prefix.as_ref(), start_key.as_ref()
+					)
+				)
+				.map(|v| v.take(count as usize).collect())
+				.map_err(client_err)))
+	}
+
 	fn storage(
 		&self,
 		block: Option<Block::Hash>,

substrate/client/rpc/src/state/state_light.rs

@@ -199,6 +199,16 @@ impl<Block, F, B, E, RA> StateBackend<B, E, Block, RA> for LightState<Block, F,
 		Box::new(result(Err(client_err(ClientError::NotAvailableOnLightClient))))
 	}
 
+	fn storage_keys_paged(
+		&self,
+		_block: Option<Block::Hash>,
+		_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>,

substrate/client/src/client.rs

@@ -99,6 +99,46 @@ pub struct Client<B, E, Block, RA> where Block: BlockT {
 	_phantom: PhantomData<RA>,
 }
 
+/// An `Iterator` that iterates keys in a given block under a prefix.
+pub struct KeyIterator<'a, State, Block> {
+	state: State,
+	prefix: Option<&'a StorageKey>,
+	current_key: Vec<u8>,
+	_phantom: PhantomData<Block>,
+}
+
+impl <'a, State, Block> KeyIterator<'a, State, Block> {
+	fn new(state: State, prefix: Option<&'a StorageKey>, current_key: Vec<u8>) -> Self {
+		Self {
+			state,
+			prefix,
+			current_key,
+			_phantom: PhantomData,
+		}
+	}
+}
+
+impl<'a, State, Block> Iterator for KeyIterator<'a, State, Block> where
+	Block: BlockT,
+	State: StateBackend<HasherFor<Block>>,
+{
+	type Item = StorageKey;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		let next_key = self.state
+			.next_storage_key(&self.current_key)
+			.ok()
+			.flatten()?;
+		if let Some(prefix) = self.prefix {
+			if !next_key.starts_with(&prefix.0[..]) {
+				return None;
+			}
+		}
+		self.current_key = next_key.clone();
+		Some(StorageKey(next_key))
+	}
+}
+
 // used in importing a block, where additional changes are made after the runtime
 // executed.
 enum PrePostHeader<H> {
@@ -234,12 +274,27 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 		self.backend.state_at(*block)
 	}
 
-	/// Given a `BlockId` and a key prefix, return the matching child storage keys in that block.
+	/// Given a `BlockId` and a key prefix, return the matching storage keys in that block.
 	pub fn storage_keys(&self, id: &BlockId<Block>, key_prefix: &StorageKey) -> sp_blockchain::Result<Vec<StorageKey>> {
 		let keys = self.state_at(id)?.keys(&key_prefix.0).into_iter().map(StorageKey).collect();
 		Ok(keys)
 	}
 
+	/// Given a `BlockId` and a key prefix, return a `KeyIterator` iterates matching storage keys in that block.
+	pub fn storage_keys_iter<'a>(
+		&self,
+		id: &BlockId<Block>,
+		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(state, prefix, start_key))
+	}
+
 	/// Given a `BlockId` and a key, return the value under the key in that block.
 	pub fn storage(&self, id: &BlockId<Block>, key: &StorageKey) -> sp_blockchain::Result<Option<StorageData>> {
 		Ok(self.state_at(id)?
@@ -1900,6 +1955,7 @@ pub(crate) mod tests {
 		sc_client_db::{Backend, DatabaseSettings, DatabaseSettingsSrc, PruningMode},
 		runtime::{self, Block, Transfer, RuntimeApi, TestAPI},
 	};
+	use hex_literal::hex;
 
 	/// Returns tuple, consisting of:
 	/// 1) test client pre-filled with blocks changing balances;
@@ -3303,4 +3359,61 @@ pub(crate) mod tests {
 			vec![(30, 0), (27, 0), (25, 0), (24, 0), (11, 0)]
 		);
 	}
+
+	#[test]
+	fn storage_keys_iter_prefix_and_start_key_works() {
+		let client = substrate_test_runtime_client::new();
+
+		let prefix = StorageKey(hex!("3a").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()]);
+
+		let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), Some(&StorageKey(hex!("3a636f6465").to_vec())))
+			.unwrap()
+			.map(|x| x.0)
+			.collect();
+		assert_eq!(res, [hex!("3a686561707061676573").to_vec()]);
+
+		let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), Some(&StorageKey(hex!("3a686561707061676573").to_vec())))
+			.unwrap()
+			.map(|x| x.0)
+			.collect();
+		assert_eq!(res, Vec::<Vec<u8>>::new());
+	}
+
+	#[test]
+	fn storage_keys_iter_works() {
+		let client = substrate_test_runtime_client::new();
+
+		let prefix = StorageKey(hex!("").to_vec());
+
+		let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), None)
+			.unwrap()
+			.take(2)
+			.map(|x| x.0)
+			.collect();
+		assert_eq!(res, [hex!("0befda6e1ca4ef40219d588a727f1271").to_vec(), hex!("3a636f6465").to_vec()]);
+
+		let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), Some(&StorageKey(hex!("3a636f6465").to_vec())))
+			.unwrap()
+			.take(3)
+			.map(|x| x.0)
+			.collect();
+		assert_eq!(res, [
+			hex!("3a686561707061676573").to_vec(),
+			hex!("6644b9b8bc315888ac8e41a7968dc2b4141a5403c58acdf70b7e8f7e07bf5081").to_vec(),
+			hex!("79c07e2b1d2e2abfd4855b936617eeff5e0621c4869aa60c02be9adcc98a0d1d").to_vec(),
+		]);
+
+		let res: Vec<_> = client.storage_keys_iter(&BlockId::Number(0), Some(&prefix), Some(&StorageKey(hex!("79c07e2b1d2e2abfd4855b936617eeff5e0621c4869aa60c02be9adcc98a0d1d").to_vec())))
+			.unwrap()
+			.take(1)
+			.map(|x| x.0)
+			.collect();
+		assert_eq!(res, [hex!("cf722c0832b5231d35e29f319ff27389f5032bfc7bfc3ba5ed7839f2042fb99f").to_vec()]);
+	}
 }