chainHead: Add support for storage pagination and cancellation

client/rpc-spec-v2/Cargo.toml

@@ -36,6 +36,7 @@ tokio = { version = "1.22.0", features = ["sync"] }
 array-bytes = "6.1"
 log = "0.4.17"
 futures-util = { version = "0.3.19", default-features = false }
+async-channel = "1.8.0"
 
 [dev-dependencies]
 serde_json = "1.0"

client/rpc-spec-v2/src/chain_head/api.rs

@@ -119,4 +119,31 @@ pub trait ChainHeadApi<Hash> {
 	/// This method is unstable and subject to change in the future.
 	#[method(name = "chainHead_unstable_unpin", blocking)]
 	fn chain_head_unstable_unpin(&self, follow_subscription: String, hash: Hash) -> RpcResult<()>;
+
+	/// Resumes a storage fetch started with `chainHead_storage` after it has generated an
+	/// `operationWaitingForContinue` event.
+	///
+	/// # Unstable
+	///
+	/// This method is unstable and subject to change in the future.
+	#[method(name = "chainHead_unstable_continue", blocking)]
+	fn chain_head_unstable_continue(
+		&self,
+		follow_subscription: String,
+		operation_id: String,
+	) -> RpcResult<()>;
+
+	/// Stops an operation started with chainHead_unstable_body, chainHead_unstable_call, or
+	/// chainHead_unstable_storage. If the operation was still in progress, this interrupts it. If
+	/// the operation was already finished, this call has no effect.
+	///
+	/// # Unstable
+	///
+	/// This method is unstable and subject to change in the future.
+	#[method(name = "chainHead_unstable_stopOperation", blocking)]
+	fn chain_head_unstable_stop_operation(
+		&self,
+		follow_subscription: String,
+		operation_id: String,
+	) -> RpcResult<()>;
 }

client/rpc-spec-v2/src/chain_head/chain_head.rs

@@ -61,6 +61,9 @@ pub struct ChainHeadConfig {
 	pub subscription_max_pinned_duration: Duration,
 	/// The maximum number of ongoing operations per subscription.
 	pub subscription_max_ongoing_operations: usize,
+	/// The maximum number of items reported by the `chainHead_storage` before
+	/// pagination is required.
+	pub operation_max_storage_items: usize,
 }
 
 /// Maximum pinned blocks across all connections.
@@ -78,12 +81,17 @@ const MAX_PINNED_DURATION: Duration = Duration::from_secs(60);
 /// Note: The lower limit imposed by the spec is 16.
 const MAX_ONGOING_OPERATIONS: usize = 16;
 
+/// The maximum number of items the `chainHead_storage` can return
+/// before paginations is required.
+const MAX_STORAGE_ITER_ITEMS: usize = 5;
+
 impl Default for ChainHeadConfig {
 	fn default() -> Self {
 		ChainHeadConfig {
 			global_max_pinned_blocks: MAX_PINNED_BLOCKS,
 			subscription_max_pinned_duration: MAX_PINNED_DURATION,
 			subscription_max_ongoing_operations: MAX_ONGOING_OPERATIONS,
+			operation_max_storage_items: MAX_STORAGE_ITER_ITEMS,
 		}
 	}
 }
@@ -100,6 +108,9 @@ pub struct ChainHead<BE: Backend<Block>, Block: BlockT, Client> {
 	subscriptions: Arc<SubscriptionManagement<Block, BE>>,
 	/// The hexadecimal encoded hash of the genesis block.
 	genesis_hash: String,
+	/// The maximum number of items reported by the `chainHead_storage` before
+	/// pagination is required.
+	operation_max_storage_items: usize,
 	/// Phantom member to pin the block type.
 	_phantom: PhantomData<Block>,
 }
@@ -124,6 +135,7 @@ impl<BE: Backend<Block>, Block: BlockT, Client> ChainHead<BE, Block, Client> {
 				config.subscription_max_ongoing_operations,
 				backend,
 			)),
+			operation_max_storage_items: config.operation_max_storage_items,
 			genesis_hash,
 			_phantom: PhantomData,
 		}
@@ -243,6 +255,7 @@ where
 			Err(_) => return Err(ChainHeadRpcError::InvalidBlock.into()),
 		};
 
+		let operation = block_guard.operation();
 		let event = match self.client.block(hash) {
 			Ok(Some(signed_block)) => {
 				let extrinsics = signed_block
@@ -252,7 +265,7 @@ where
 					.map(|extrinsic| hex_string(&extrinsic.encode()))
 					.collect();
 				FollowEvent::<Block::Hash>::OperationBodyDone(OperationBodyDone {
-					operation_id: block_guard.operation_id(),
+					operation_id: operation.operation_id(),
 					value: extrinsics,
 				})
 			},
@@ -268,14 +281,14 @@ where
 				return Err(ChainHeadRpcError::InvalidBlock.into())
 			},
 			Err(error) => FollowEvent::<Block::Hash>::OperationError(OperationError {
-				operation_id: block_guard.operation_id(),
+				operation_id: operation.operation_id(),
 				error: error.to_string(),
 			}),
 		};
 
 		let _ = block_guard.response_sender().unbounded_send(event);
 		Ok(MethodResponse::Started(MethodResponseStarted {
-			operation_id: block_guard.operation_id(),
+			operation_id: operation.operation_id(),
 			discarded_items: None,
 		}))
 	}
@@ -349,17 +362,21 @@ where
 				Err(_) => return Err(ChainHeadRpcError::InvalidBlock.into()),
 			};
 
-		let storage_client = ChainHeadStorage::<Client, Block, BE>::new(self.client.clone());
-		let operation_id = block_guard.operation_id();
+		let mut storage_client = ChainHeadStorage::<Client, Block, BE>::new(
+			self.client.clone(),
+			self.operation_max_storage_items,
+		);
+		let operation = block_guard.operation();
+		let operation_id = operation.operation_id();
 
 		// The number of operations we are allowed to execute.
-		let num_operations = block_guard.num_reserved();
+		let num_operations = operation.num_reserved();
 		let discarded = items.len().saturating_sub(num_operations);
 		let mut items = items;
 		items.truncate(num_operations);
 
 		let fut = async move {
-			storage_client.generate_events(block_guard, hash, items, child_trie);
+			storage_client.generate_events(block_guard, hash, items, child_trie).await;
 		};
 
 		self.executor
@@ -401,26 +418,27 @@ where
 			.into())
 		}
 
+		let operation = block_guard.operation();
 		let event = self
 			.client
 			.executor()
 			.call(hash, &function, &call_parameters, CallContext::Offchain)
 			.map(|result| {
 				FollowEvent::<Block::Hash>::OperationCallDone(OperationCallDone {
-					operation_id: block_guard.operation_id(),
+					operation_id: operation.operation_id(),
 					output: hex_string(&result),
 				})
 			})
 			.unwrap_or_else(|error| {
 				FollowEvent::<Block::Hash>::OperationError(OperationError {
-					operation_id: block_guard.operation_id(),
+					operation_id: operation.operation_id(),
 					error: error.to_string(),
 				})
 			});
 
 		let _ = block_guard.response_sender().unbounded_send(event);
 		Ok(MethodResponse::Started(MethodResponseStarted {
-			operation_id: block_guard.operation_id(),
+			operation_id: operation.operation_id(),
 			discarded_items: None,
 		}))
 	}
@@ -443,4 +461,35 @@ where
 			Err(_) => Err(ChainHeadRpcError::InvalidBlock.into()),
 		}
 	}
+
+	fn chain_head_unstable_continue(
+		&self,
+		follow_subscription: String,
+		operation_id: String,
+	) -> RpcResult<()> {
+		let Some(operation) = self.subscriptions.get_operation(&follow_subscription, &operation_id) else {
+			return Ok(())
+		};
+
+		if !operation.submit_continue() {
+			// Continue called without generating a `WaitingForContinue` event.
+			Err(ChainHeadRpcError::InvalidContinue.into())
+		} else {
+			Ok(())
+		}
+	}
+
+	fn chain_head_unstable_stop_operation(
+		&self,
+		follow_subscription: String,
+		operation_id: String,
+	) -> RpcResult<()> {
+		let Some(operation) = self.subscriptions.get_operation(&follow_subscription, &operation_id) else {
+			return Ok(())
+		};
+
+		operation.stop_operation();
+
+		Ok(())
+	}
 }