archive: Refactor archive_storage method into subscription

prdoc/pr_5997.prdoc

@@ -0,0 +1,18 @@
+# Schema: Polkadot SDK PRDoc Schema (prdoc) v1.0.0
+# See doc at https://raw.githubusercontent.com/paritytech/polkadot-sdk/master/prdoc/schema_user.json
+
+title: Implement archive_unstable_storageDiff method
+
+doc:
+  - audience: Node Dev
+    description: |
+      This PR implements the `archive_unstable_storageDiff` rpc-v2 method.
+      Developers can use this method to fetch the storage differences
+      between two blocks. This is useful for oracles and archive nodes.
+      For more details see: https://github.com/paritytech/json-rpc-interface-spec/blob/main/src/api/archive_unstable_storageDiff.md.
+
+crates:
+  - name: sc-rpc-spec-v2
+    bump: major
+  - name: sc-service
+    bump: patch

substrate/client/rpc-spec-v2/Cargo.toml

@@ -42,6 +42,7 @@ log = { workspace = true, default-features = true }
 futures-util = { workspace = true }
 rand = { workspace = true, default-features = true }
 schnellru = { workspace = true }
+itertools = { workspace = true }
 
 [dev-dependencies]
 jsonrpsee = { workspace = true, features = ["server", "ws-client"] }

substrate/client/rpc-spec-v2/src/archive/api.rs

@@ -19,7 +19,9 @@
 //! API trait of the archive methods.
 
 use crate::{
-	common::events::{ArchiveStorageResult, PaginatedStorageQuery},
+	common::events::{
+		ArchiveStorageDiffEvent, ArchiveStorageDiffItem, ArchiveStorageEvent, StorageQuery,
+	},
 	MethodResult,
 };
 use jsonrpsee::{core::RpcResult, proc_macros::rpc};
@@ -97,11 +99,32 @@ pub trait ArchiveApi<Hash> {
 	/// # Unstable
 	///
 	/// This method is unstable and subject to change in the future.
-	#[method(name = "archive_unstable_storage", blocking)]
+	#[subscription(
+		name = "archive_unstable_storage" => "archive_unstable_storageEvent",
+		unsubscribe = "archive_unstable_stopStorage",
+		item = ArchiveStorageEvent,
+	)]
 	fn archive_unstable_storage(
 		&self,
 		hash: Hash,
-		items: Vec<PaginatedStorageQuery<String>>,
+		items: Vec<StorageQuery<String>>,
 		child_trie: Option<String>,
-	) -> RpcResult<ArchiveStorageResult>;
+	);
+
+	/// Returns the storage difference between two blocks.
+	///
+	/// # Unstable
+	///
+	/// This method is unstable and subject to change in the future.
+	#[subscription(
+		name = "archive_unstable_storageDiff" => "archive_unstable_storageDiffEvent",
+		unsubscribe = "archive_unstable_storageDiff_stopStorageDiff",
+		item = ArchiveStorageDiffEvent,
+	)]
+	fn archive_unstable_storage_diff(
+		&self,
+		hash: Hash,
+		items: Vec<ArchiveStorageDiffItem<String>>,
+		previous_hash: Option<Hash>,
+	);
 }

substrate/client/rpc-spec-v2/src/archive/archive.rs

@@ -19,17 +19,31 @@
 //! API implementation for `archive`.
 
 use crate::{
-	archive::{error::Error as ArchiveError, ArchiveApiServer},
-	common::events::{ArchiveStorageResult, PaginatedStorageQuery},
-	hex_string, MethodResult,
+	archive::{
+		archive_storage::{deduplicate_storage_diff_items, ArchiveStorageDiff},
+		error::Error as ArchiveError,
+		ArchiveApiServer,
+	},
+	common::{
+		events::{
+			ArchiveStorageDiffEvent, ArchiveStorageDiffItem, ArchiveStorageEvent, StorageQuery,
+		},
+		storage::{QueryResult, StorageSubscriptionClient},
+	},
+	hex_string, MethodResult, SubscriptionTaskExecutor,
 };
 
 use codec::Encode;
-use jsonrpsee::core::{async_trait, RpcResult};
+use futures::FutureExt;
+use jsonrpsee::{
+	core::{async_trait, RpcResult},
+	PendingSubscriptionSink,
+};
 use sc_client_api::{
 	Backend, BlockBackend, BlockchainEvents, CallExecutor, ChildInfo, ExecutorProvider, StorageKey,
 	StorageProvider,
 };
+use sc_rpc::utils::Subscription;
 use sp_api::{CallApiAt, CallContext};
 use sp_blockchain::{
 	Backend as BlockChainBackend, Error as BlockChainError, HeaderBackend, HeaderMetadata,
@@ -41,51 +55,26 @@ use sp_runtime::{
 };
 use std::{collections::HashSet, marker::PhantomData, sync::Arc};
 
-use super::archive_storage::ArchiveStorage;
-
-/// The configuration of [`Archive`].
-pub struct ArchiveConfig {
-	/// The maximum number of items the `archive_storage` can return for a descendant query before
-	/// pagination is required.
-	pub max_descendant_responses: usize,
-	/// The maximum number of queried items allowed for the `archive_storage` at a time.
-	pub max_queried_items: usize,
-}
+use tokio::sync::mpsc;
 
-/// The maximum number of items the `archive_storage` can return for a descendant query before
-/// pagination is required.
-///
-/// Note: this is identical to the `chainHead` value.
-const MAX_DESCENDANT_RESPONSES: usize = 5;
+pub(crate) const LOG_TARGET: &str = "rpc-spec-v2::archive";
 
-/// The maximum number of queried items allowed for the `archive_storage` at a time.
+/// The buffer capacity for each storage query.
 ///
-/// Note: A queried item can also be a descendant query which can return up to
-/// `MAX_DESCENDANT_RESPONSES`.
-const MAX_QUERIED_ITEMS: usize = 8;
-
-impl Default for ArchiveConfig {
-	fn default() -> Self {
-		Self {
-			max_descendant_responses: MAX_DESCENDANT_RESPONSES,
-			max_queried_items: MAX_QUERIED_ITEMS,
-		}
-	}
-}
+/// This is small because the underlying JSON-RPC server has
+/// its down buffer capacity per connection as well.
+const STORAGE_QUERY_BUF: usize = 16;
 
 /// An API for archive RPC calls.
 pub struct Archive<BE: Backend<Block>, Block: BlockT, Client> {
 	/// Substrate client.
 	client: Arc<Client>,
 	/// Backend of the chain.
 	backend: Arc<BE>,
+	/// Executor to spawn subscriptions.
+	executor: SubscriptionTaskExecutor,
 	/// The hexadecimal encoded hash of the genesis block.
 	genesis_hash: String,
-	/// The maximum number of items the `archive_storage` can return for a descendant query before
-	/// pagination is required.
-	storage_max_descendant_responses: usize,
-	/// The maximum number of queried items allowed for the `archive_storage` at a time.
-	storage_max_queried_items: usize,
 	/// Phantom member to pin the block type.
 	_phantom: PhantomData<Block>,
 }
@@ -96,17 +85,10 @@ impl<BE: Backend<Block>, Block: BlockT, Client> Archive<BE, Block, Client> {
 		client: Arc<Client>,
 		backend: Arc<BE>,
 		genesis_hash: GenesisHash,
-		config: ArchiveConfig,
+		executor: SubscriptionTaskExecutor,
 	) -> Self {
 		let genesis_hash = hex_string(&genesis_hash.as_ref());
-		Self {
-			client,
-			backend,
-			genesis_hash,
-			storage_max_descendant_responses: config.max_descendant_responses,
-			storage_max_queried_items: config.max_queried_items,
-			_phantom: PhantomData,
-		}
+		Self { client, backend, executor, genesis_hash, _phantom: PhantomData }
 	}
 }
 
@@ -236,46 +218,142 @@ where
 
 	fn archive_unstable_storage(
 		&self,
+		pending: PendingSubscriptionSink,
 		hash: Block::Hash,
-		items: Vec<PaginatedStorageQuery<String>>,
+		items: Vec<StorageQuery<String>>,
 		child_trie: Option<String>,
-	) -> RpcResult<ArchiveStorageResult> {
-		let items = items
-			.into_iter()
-			.map(|query| {
-				let key = StorageKey(parse_hex_param(query.key)?);
-				let pagination_start_key = query
-					.pagination_start_key
-					.map(|key| parse_hex_param(key).map(|key| StorageKey(key)))
-					.transpose()?;
-
-				// Paginated start key is only supported
-				if pagination_start_key.is_some() && !query.query_type.is_descendant_query() {
-					return Err(ArchiveError::InvalidParam(
-						"Pagination start key is only supported for descendants queries"
-							.to_string(),
-					))
-				}
-
-				Ok(PaginatedStorageQuery {
-					key,
-					query_type: query.query_type,
-					pagination_start_key,
+	) {
+		let mut storage_client =
+			StorageSubscriptionClient::<Client, Block, BE>::new(self.client.clone());
+
+		let fut = async move {
+			let Ok(mut sink) = pending.accept().await.map(Subscription::from) else { return };
+
+			let items = match items
+				.into_iter()
+				.map(|query| {
+					let key = StorageKey(parse_hex_param(query.key)?);
+					Ok(StorageQuery { key, query_type: query.query_type })
 				})
-			})
-			.collect::<Result<Vec<_>, ArchiveError>>()?;
+				.collect::<Result<Vec<_>, ArchiveError>>()
+			{
+				Ok(items) => items,
+				Err(error) => {
+					let _ = sink.send(&ArchiveStorageEvent::err(error.to_string()));
+					return
+				},
+			};
+
+			let child_trie = child_trie.map(|child_trie| parse_hex_param(child_trie)).transpose();
+			let child_trie = match child_trie {
+				Ok(child_trie) => child_trie.map(ChildInfo::new_default_from_vec),
+				Err(error) => {
+					let _ = sink.send(&ArchiveStorageEvent::err(error.to_string()));
+					return
+				},
+			};
+
+			let (tx, mut rx) = tokio::sync::mpsc::channel(STORAGE_QUERY_BUF);
+			let storage_fut = storage_client.generate_events(hash, items, child_trie, tx);
+
+			// We don't care about the return value of this join:
+			// - process_events might encounter an error (if the client disconnected)
+			// - storage_fut might encounter an error while processing a trie queries and
+			// the error is propagated via the sink.
+			let _ = futures::future::join(storage_fut, process_storage_events(&mut rx, &mut sink))
+				.await;
+		};
+
+		self.executor.spawn("substrate-rpc-subscription", Some("rpc"), fut.boxed());
+	}
 
-		let child_trie = child_trie
-			.map(|child_trie| parse_hex_param(child_trie))
-			.transpose()?
-			.map(ChildInfo::new_default_from_vec);
+	fn archive_unstable_storage_diff(
+		&self,
+		pending: PendingSubscriptionSink,
+		hash: Block::Hash,
+		items: Vec<ArchiveStorageDiffItem<String>>,
+		previous_hash: Option<Block::Hash>,
+	) {
+		let storage_client = ArchiveStorageDiff::new(self.client.clone());
+		let client = self.client.clone();
+
+		log::trace!(target: LOG_TARGET, "Storage diff subscription started");
+
+		let fut = async move {
+			let Ok(mut sink) = pending.accept().await.map(Subscription::from) else { return };
+
+			// Deduplicate the items.
+			let mut trie_items = match deduplicate_storage_diff_items(items) {
+				Ok(items) => items,
+				Err(error) => {
+					let _ = sink.send(&ArchiveStorageDiffEvent::err(error.to_string())).await;
+					return
+				},
+			};
+			// Default to using the main storage trie if no items are provided.
+			if trie_items.is_empty() {
+				trie_items.push(Vec::new());
+			}
+			log::trace!(target: LOG_TARGET, "Storage diff deduplicated items: {:?}", trie_items);
+
+			let previous_hash = if let Some(previous_hash) = previous_hash {
+				previous_hash
+			} else {
+				let Ok(Some(current_header)) = client.header(hash) else {
+					let message = format!("Block header is not present: {hash}");
+					let _ = sink.send(&ArchiveStorageDiffEvent::err(message)).await;
+					return
+				};
+				*current_header.parent_hash()
+			};
+
+			let (tx, mut rx) = tokio::sync::mpsc::channel(STORAGE_QUERY_BUF);
+			let storage_fut =
+				storage_client.handle_trie_queries(hash, previous_hash, trie_items, tx.clone());
+
+			// We don't care about the return value of this join:
+			// - process_events might encounter an error (if the client disconnected)
+			// - storage_fut might encounter an error while processing a trie queries and
+			// the error is propagated via the sink.
+			let _ = futures::future::join(storage_fut, process_events(&mut rx, &mut sink)).await;
+		};
+
+		self.executor.spawn("substrate-rpc-subscription", Some("rpc"), fut.boxed());
+	}
+}
+
+/// Sends all the events to the sink.
+async fn process_events(rx: &mut mpsc::Receiver<ArchiveStorageDiffEvent>, sink: &mut Subscription) {
+	while let Some(event) = rx.recv().await {
+		if event.is_done() {
+			log::debug!(target: LOG_TARGET, "Finished processing partial trie query");
+		} else if event.is_err() {
+			log::debug!(target: LOG_TARGET, "Error encountered while processing partial trie query");
+		}
 
-		let storage_client = ArchiveStorage::new(
-			self.client.clone(),
-			self.storage_max_descendant_responses,
-			self.storage_max_queried_items,
-		);
+		if sink.send(&event).await.is_err() {
+			return
+		}
+	}
+}
 
-		Ok(storage_client.handle_query(hash, items, child_trie))
+/// Sends all the events to the sink.
+async fn process_storage_events(rx: &mut mpsc::Receiver<QueryResult>, sink: &mut Subscription) {
+	while let Some(event) = rx.recv().await {
+		match event {
+			Ok(None) => continue,
+
+			Ok(Some(event)) =>
+				if sink.send(&ArchiveStorageEvent::result(event)).await.is_err() {
+					return
+				},
+
+			Err(error) =>
+				if sink.send(&ArchiveStorageEvent::err(error)).await.is_err() {
+					return
+				},
+		}
 	}
+
+	let _ = sink.send(&ArchiveStorageEvent::StorageDone).await;
 }