Add storage query functions for multiple keys

fixes paritytech#9203

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

@@ -66,6 +66,15 @@ pub trait ChildStateApi<Hash> {
 		hash: Option<Hash>,
 	) -> FutureResult<Option<StorageData>>;
 
+	/// Returns child storage entries for multiple keys at a specific block's state.
+	#[rpc(name = "childstate_getStorages")]
+	fn storages(
+		&self,
+		child_storage_key: PrefixedStorageKey,
+		keys: Vec<StorageKey>,
+		hash: Option<Hash>,
+	) -> FutureResult<Vec<Option<StorageData>>>;
+
 	/// Returns the hash of a child storage entry at a block's state.
 	#[rpc(name = "childstate_getStorageHash")]
 	fn storage_hash(

client/rpc/src/state/mod.rs

@@ -96,6 +96,13 @@ where
 		key: StorageKey,
 	) -> FutureResult<Option<StorageData>>;
 
+	/// Returns storage entries for multiple keys at a specific block's state.
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>>;
+
 	/// Returns the hash of a storage entry at a block's state.
 	fn storage_hash(
 		&self,
@@ -462,6 +469,14 @@ where
 		key: StorageKey,
 	) -> FutureResult<Option<StorageData>>;
 
+	/// Returns child storage entries at a specific block's state.
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		storage_key: PrefixedStorageKey,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>>;
+
 	/// Returns the hash of a child storage entry at a block's state.
 	fn storage_hash(
 		&self,
@@ -511,6 +526,15 @@ where
 		self.backend.storage(block, storage_key, key)
 	}
 
+	fn storages(
+		&self,
+		storage_key: PrefixedStorageKey,
+		keys: Vec<StorageKey>,
+		block: Option<Block::Hash>,
+	) -> FutureResult<Vec<Option<StorageData>>> {
+		self.backend.storages(block, storage_key, keys)
+	}
+
 	fn storage_keys(
 		&self,
 		storage_key: PrefixedStorageKey,

client/rpc/src/state/state_full.rs

@@ -22,7 +22,10 @@ use futures::{future, StreamExt as _, TryStreamExt as _};
 use jsonrpc_pubsub::{manager::SubscriptionManager, typed::Subscriber, SubscriptionId};
 use log::warn;
 use rpc::{
-	futures::{future::result, stream, Future, Sink, Stream},
+	futures::{
+		future::{join_all, result},
+		stream, Future, Sink, Stream,
+	},
 	Result as RpcResult,
 };
 use std::{
@@ -358,6 +361,22 @@ where
 		))
 	}
 
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>> {
+		let block = match self.block_or_best(block) {
+			Ok(b) => b,
+			Err(e) => return Box::new(result(Err(client_err(e)))),
+		};
+		let client = self.client.clone();
+		Box::new(join_all(
+			keys.into_iter()
+				.map(move |key| client.storage(&BlockId::Hash(block), &key).map_err(client_err)),
+		))
+	}
+
 	fn storage_size(
 		&self,
 		block: Option<Block::Hash>,
@@ -722,6 +741,30 @@ where
 		))
 	}
 
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		storage_key: PrefixedStorageKey,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>> {
+		let block = match self.block_or_best(block) {
+			Ok(b) => b,
+			Err(e) => return Box::new(result(Err(client_err(e)))),
+		};
+		let client = self.client.clone();
+		Box::new(
+			join_all(keys.into_iter().map(move |key| {
+				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),
+				};
+				client.child_storage(&BlockId::Hash(block), &child_info, &key)
+			}))
+			.map_err(client_err),
+		)
+	}
+
 	fn storage_hash(
 		&self,
 		block: Option<Block::Hash>,

client/rpc/src/state/state_light.rs

@@ -30,7 +30,7 @@ use log::warn;
 use parking_lot::Mutex;
 use rpc::{
 	futures::{
-		future::{result, Future},
+		future::{join_all, result, Future},
 		stream::Stream,
 		Sink,
 	},
@@ -245,6 +245,26 @@ where
 		)
 	}
 
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>> {
+		let remote_blockchain = self.remote_blockchain.clone();
+		let fetcher = self.fetcher.clone();
+		let block = self.block_or_best(block);
+		Box::new(join_all(keys.into_iter().map(move |key| {
+			storage(&*remote_blockchain, fetcher.clone(), block, vec![key.0.clone()])
+				.boxed()
+				.compat()
+				.map(move |mut values| {
+					values
+						.remove(&key)
+						.expect("successful request has entries for all requested keys; qed")
+				})
+		})))
+	}
+
 	fn storage_hash(
 		&self,
 		block: Option<Block::Hash>,
@@ -563,6 +583,52 @@ where
 		Box::new(child_storage.boxed().compat())
 	}
 
+	fn storages(
+		&self,
+		block: Option<Block::Hash>,
+		storage_key: PrefixedStorageKey,
+		keys: Vec<StorageKey>,
+	) -> FutureResult<Vec<Option<StorageData>>> {
+		let block = self.block_or_best(block);
+		let fetcher = self.fetcher.clone();
+		let remote_blockchain = self.remote_blockchain.clone();
+		Box::new(join_all(keys.into_iter().map(move |key| {
+			let storage_key = storage_key.clone();
+			let fetcher2 = fetcher.clone();
+			let child_storage =
+				resolve_header(&*remote_blockchain, &*fetcher, block).then(move |result| {
+					match result {
+						Ok(header) => Either::Left(
+							fetcher2
+								.remote_read_child(RemoteReadChildRequest {
+									block,
+									header,
+									storage_key,
+									keys: vec![key.0.clone()],
+									retry_count: Default::default(),
+								})
+								.then(move |result| {
+									ready(
+										result
+											.map(|mut data| {
+												data.remove(&key.0)
+												.expect(
+													"successful result has entry for all keys; qed",
+												)
+												.map(StorageData)
+											})
+											.map_err(client_err),
+									)
+								}),
+						),
+						Err(error) => Either::Right(ready(Err(error))),
+					}
+				});
+
+			Box::new(child_storage.boxed().compat())
+		})))
+	}
+
 	fn storage_hash(
 		&self,
 		block: Option<Block::Hash>,

client/rpc/src/state/tests.rs

@@ -98,6 +98,41 @@ fn should_return_storage() {
 	);
 }
 
+#[test]
+fn should_return_storages() {
+	const KEY1: &[u8] = b":mock";
+	const KEY2: &[u8] = b":turtle";
+	const VALUE: &[u8] = b"hello world";
+	const CHILD_VALUE1: &[u8] = b"hello world !";
+	const CHILD_VALUE2: &[u8] = b"hello world    !";
+
+	let child_info = ChildInfo::new_default(STORAGE_KEY);
+	let client = TestClientBuilder::new()
+		.add_extra_storage(KEY1.to_vec(), VALUE.to_vec())
+		.add_extra_child_storage(&child_info, KEY1.to_vec(), CHILD_VALUE1.to_vec())
+		.add_extra_child_storage(&child_info, KEY2.to_vec(), CHILD_VALUE2.to_vec())
+		.build();
+	let genesis_hash = client.genesis_hash();
+	let (_client, child) = new_full(
+		Arc::new(client),
+		SubscriptionManager::new(Arc::new(TaskExecutor)),
+		DenyUnsafe::No,
+		None,
+	);
+	let keys = &[StorageKey(KEY1.to_vec()), StorageKey(KEY2.to_vec())];
+
+	assert_eq!(
+		executor::block_on(
+			child
+				.storages(prefixed_storage_key(), keys.to_vec(), Some(genesis_hash).into())
+				.map(|x| x.into_iter().map(|x| x.map(|y| y.0.len()).unwrap()).sum::<usize>())
+				.compat(),
+		)
+		.unwrap(),
+		CHILD_VALUE1.len() + CHILD_VALUE2.len(),
+	);
+}
+
 #[test]
 fn should_return_child_storage() {
 	let child_info = ChildInfo::new_default(STORAGE_KEY);
@@ -130,6 +165,49 @@ fn should_return_child_storage() {
 	assert_matches!(child.storage_size(child_key.clone(), key.clone(), None).wait(), Ok(Some(1)));
 }
 
+#[test]
+fn should_return_child_storages() {
+	let child_info = ChildInfo::new_default(STORAGE_KEY);
+	let client = Arc::new(
+		substrate_test_runtime_client::TestClientBuilder::new()
+			.add_child_storage(&child_info, "key1", vec![42_u8])
+			.add_child_storage(&child_info, "key2", vec![43_u8, 44])
+			.build(),
+	);
+	let genesis_hash = client.genesis_hash();
+	let (_client, child) =
+		new_full(client, SubscriptionManager::new(Arc::new(TaskExecutor)), DenyUnsafe::No, None);
+	let child_key = prefixed_storage_key();
+	let keys = vec![StorageKey(b"key1".to_vec()), StorageKey(b"key2".to_vec())];
+
+	let res = child
+		.storages(child_key.clone(), keys.clone(), Some(genesis_hash).into())
+		.wait()
+		.unwrap();
+
+	assert_matches!(
+		res[0],
+		Some(StorageData(ref d))
+			if d[0] == 42 && d.len() == 1
+	);
+	assert_matches!(
+		res[1],
+		Some(StorageData(ref d))
+			if d[0] == 43 && d[1] == 44 && d.len() == 2
+	);
+	assert_matches!(
+		child
+			.storage_hash(child_key.clone(), keys[0].clone(), Some(genesis_hash).into(),)
+			.wait()
+			.map(|x| x.is_some()),
+		Ok(true)
+	);
+	assert_matches!(
+		child.storage_size(child_key.clone(), keys[0].clone(), None).wait(),
+		Ok(Some(1))
+	);
+}
+
 #[test]
 fn should_call_contract() {
 	let client = Arc::new(substrate_test_runtime_client::new());