chainHead: Allow methods to be called from within a single connection context and limit connections (#3481)

This PR ensures that the chainHead RPC class can be called only from within the same connection context. The chainHead methods are now registered as raw methods. - https://github.com/paritytech/jsonrpsee/pull/1297 The concept of raw methods is introduced in jsonrpsee, which is an async method that exposes the connection ID: The raw method doesn't have the concept of a blocking method. Previously blocking methods are now spawning a blocking task to handle their blocking (ie DB) access. We spawn the same number of tasks as before, however we do that explicitly. Another approach would be implementing a RPC middleware that captures and decodes the method parameters: - https://github.com/paritytech/polkadot-sdk/pull/3343 However, that approach is prone to errors since the methods are hardcoded by name. Performace is affected by the double deserialization that needs to happen to extract the subscription ID we'd like to limit. Once from the middleware, and once from the methods itself. This PR paves the way to implement the chainHead connection limiter: - https://github.com/paritytech/polkadot-sdk/issues/1505 Registering tokens (subscription ID / operation ID) on the `RpcConnections` could be extended to return an error when the maximum number of operations is reached. While at it, have added an integration-test to ensure that chainHead methods can be called from within the same connection context. Before this is merged, a new JsonRPC release should be made to expose the `raw-methods`: - [x] Use jsonrpsee from crates io (blocked by: https://github.com/paritytech/jsonrpsee/pull/1297) Closes: https://github.com/paritytech/polkadot-sdk/issues/3207 cc @paritytech/subxt-team --------- Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io> Co-authored-by: Niklas Adolfsson <niklasadolfsson1@gmail.com>
2026-07-21 07:25:44 +00:00 · 2024-04-02 16:12:34 +03:00
parent 5eff3f94be
commit 7430f41350
11 changed files with 934 additions and 154 deletions
@@ -34,10 +34,10 @@ use crate::{
 	hex_string, SubscriptionTaskExecutor,
 };
 use codec::Encode;
-use futures::future::FutureExt;
+use futures::{channel::oneshot, future::FutureExt};
 use jsonrpsee::{
-	core::async_trait, server::ResponsePayload, types::SubscriptionId, MethodResponseFuture,
-	PendingSubscriptionSink, SubscriptionSink,
+	core::async_trait, server::ResponsePayload, types::SubscriptionId, ConnectionDetails,
+	MethodResponseFuture, PendingSubscriptionSink, SubscriptionSink,
 };
 use log::debug;
 use sc_client_api::{
@@ -65,6 +65,8 @@ pub struct ChainHeadConfig {
 	/// The maximum number of items reported by the `chainHead_storage` before
 	/// pagination is required.
 	pub operation_max_storage_items: usize,
+	/// The maximum number of `chainHead_follow` subscriptions per connection.
+	pub max_follow_subscriptions_per_connection: usize,
 }

 /// Maximum pinned blocks across all connections.
@@ -86,6 +88,9 @@ const MAX_ONGOING_OPERATIONS: usize = 16;
 /// before paginations is required.
 const MAX_STORAGE_ITER_ITEMS: usize = 5;

+/// The maximum number of `chainHead_follow` subscriptions per connection.
+const MAX_FOLLOW_SUBSCRIPTIONS_PER_CONNECTION: usize = 4;
+
 impl Default for ChainHeadConfig {
 	fn default() -> Self {
 		ChainHeadConfig {
@@ -93,6 +98,7 @@ impl Default for ChainHeadConfig {
 			subscription_max_pinned_duration: MAX_PINNED_DURATION,
 			subscription_max_ongoing_operations: MAX_ONGOING_OPERATIONS,
 			operation_max_storage_items: MAX_STORAGE_ITER_ITEMS,
+			max_follow_subscriptions_per_connection: MAX_FOLLOW_SUBSCRIPTIONS_PER_CONNECTION,
 		}
 	}
 }
@@ -106,7 +112,7 @@ pub struct ChainHead<BE: Backend<Block>, Block: BlockT, Client> {
 	/// Executor to spawn subscriptions.
 	executor: SubscriptionTaskExecutor,
 	/// Keep track of the pinned blocks for each subscription.
-	subscriptions: Arc<SubscriptionManagement<Block, BE>>,
+	subscriptions: SubscriptionManagement<Block, BE>,
 	/// The maximum number of items reported by the `chainHead_storage` before
 	/// pagination is required.
 	operation_max_storage_items: usize,
@@ -126,12 +132,13 @@ impl<BE: Backend<Block>, Block: BlockT, Client> ChainHead<BE, Block, Client> {
 			client,
 			backend: backend.clone(),
 			executor,
-			subscriptions: Arc::new(SubscriptionManagement::new(
+			subscriptions: SubscriptionManagement::new(
 				config.global_max_pinned_blocks,
 				config.subscription_max_pinned_duration,
 				config.subscription_max_ongoing_operations,
+				config.max_follow_subscriptions_per_connection,
 				backend,
-			)),
+			),
 			operation_max_storage_items: config.operation_max_storage_items,
 			_phantom: PhantomData,
 		}
@@ -182,12 +189,23 @@ where
 		let client = self.client.clone();

 		let fut = async move {
+			// Ensure the current connection ID has enough space to accept a new subscription.
+			let connection_id = pending.connection_id();
+			// The RAII `reserved_subscription` will clean up resources on drop:
+			// - free the reserved subscription for the connection ID.
+			// - remove the subscription ID from the subscription management.
+			let Some(mut reserved_subscription) = subscriptions.reserve_subscription(connection_id)
+			else {
+				pending.reject(ChainHeadRpcError::ReachedLimits).await;
+				return
+			};
+
 			let Ok(sink) = pending.accept().await else { return };

 			let sub_id = read_subscription_id_as_string(&sink);
-
 			// Keep track of the subscription.
-			let Some(sub_data) = subscriptions.insert_subscription(sub_id.clone(), with_runtime)
+			let Some(sub_data) =
+				reserved_subscription.insert_subscription(sub_id.clone(), with_runtime)
 			else {
 				// Inserting the subscription can only fail if the JsonRPSee
 				// generated a duplicate subscription ID.
@@ -201,91 +219,117 @@ where
 			let mut chain_head_follow = ChainHeadFollower::new(
 				client,
 				backend,
-				subscriptions.clone(),
+				subscriptions,
 				with_runtime,
 				sub_id.clone(),
 			);

 			chain_head_follow.generate_events(sink, sub_data).await;

-			subscriptions.remove_subscription(&sub_id);
 			debug!(target: LOG_TARGET, "[follow][id={:?}] Subscription removed", sub_id);
 		};

 		self.executor.spawn("substrate-rpc-subscription", Some("rpc"), fut.boxed());
 	}

-	fn chain_head_unstable_body(
+	async fn chain_head_unstable_body(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		hash: Block::Hash,
 	) -> ResponsePayload<'static, MethodResponse> {
-		let mut block_guard = match self.subscriptions.lock_block(&follow_subscription, hash, 1) {
-			Ok(block) => block,
-			Err(SubscriptionManagementError::SubscriptionAbsent) |
-			Err(SubscriptionManagementError::ExceededLimits) =>
-				return ResponsePayload::success(MethodResponse::LimitReached),
-			Err(SubscriptionManagementError::BlockHashAbsent) => {
-				// Block is not part of the subscription.
-				return ResponsePayload::error(ChainHeadRpcError::InvalidBlock);
-			},
-			Err(_) => return ResponsePayload::error(ChainHeadRpcError::InvalidBlock),
-		};
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			// The spec says to return `LimitReached` if the follow subscription is invalid or
+			// stale.
+			return ResponsePayload::success(MethodResponse::LimitReached);
+		}

-		let operation_id = block_guard.operation().operation_id();
+		let client = self.client.clone();
+		let subscriptions = self.subscriptions.clone();
+		let executor = self.executor.clone();

-		let event = match self.client.block(hash) {
-			Ok(Some(signed_block)) => {
-				let extrinsics = signed_block
-					.block
-					.extrinsics()
-					.iter()
-					.map(|extrinsic| hex_string(&extrinsic.encode()))
-					.collect();
-				FollowEvent::<Block::Hash>::OperationBodyDone(OperationBodyDone {
+		let result = spawn_blocking(&self.executor, async move {
+			let mut block_guard = match subscriptions.lock_block(&follow_subscription, hash, 1) {
+				Ok(block) => block,
+				Err(SubscriptionManagementError::SubscriptionAbsent) |
+				Err(SubscriptionManagementError::ExceededLimits) =>
+					return ResponsePayload::success(MethodResponse::LimitReached),
+				Err(SubscriptionManagementError::BlockHashAbsent) => {
+					// Block is not part of the subscription.
+					return ResponsePayload::error(ChainHeadRpcError::InvalidBlock);
+				},
+				Err(_) => return ResponsePayload::error(ChainHeadRpcError::InvalidBlock),
+			};
+
+			let operation_id = block_guard.operation().operation_id();
+
+			let event = match client.block(hash) {
+				Ok(Some(signed_block)) => {
+					let extrinsics = signed_block
+						.block
+						.extrinsics()
+						.iter()
+						.map(|extrinsic| hex_string(&extrinsic.encode()))
+						.collect();
+					FollowEvent::<Block::Hash>::OperationBodyDone(OperationBodyDone {
+						operation_id: operation_id.clone(),
+						value: extrinsics,
+					})
+				},
+				Ok(None) => {
+					// The block's body was pruned. This subscription ID has become invalid.
+					debug!(
+						target: LOG_TARGET,
+						"[body][id={:?}] Stopping subscription because hash={:?} was pruned",
+						&follow_subscription,
+						hash
+					);
+					subscriptions.remove_subscription(&follow_subscription);
+					return ResponsePayload::error(ChainHeadRpcError::InvalidBlock)
+				},
+				Err(error) => FollowEvent::<Block::Hash>::OperationError(OperationError {
 					operation_id: operation_id.clone(),
-					value: extrinsics,
-				})
-			},
-			Ok(None) => {
-				// The block's body was pruned. This subscription ID has become invalid.
-				debug!(
-					target: LOG_TARGET,
-					"[body][id={:?}] Stopping subscription because hash={:?} was pruned",
-					&follow_subscription,
-					hash
-				);
-				self.subscriptions.remove_subscription(&follow_subscription);
-				return ResponsePayload::error(ChainHeadRpcError::InvalidBlock)
-			},
-			Err(error) => FollowEvent::<Block::Hash>::OperationError(OperationError {
-				operation_id: operation_id.clone(),
-				error: error.to_string(),
-			}),
-		};
+					error: error.to_string(),
+				}),
+			};

-		let (rp, rp_fut) = method_started_response(operation_id, None);
+			let (rp, rp_fut) = method_started_response(operation_id, None);
+			let fut = async move {
+				// Wait for the server to send out the response and if it produces an error no event
+				// should be generated.
+				if rp_fut.await.is_err() {
+					return;
+				}

-		let fut = async move {
-			// Events should only by generated
-			// if the response was successfully propagated.
-			if rp_fut.await.is_err() {
-				return;
-			}
-			let _ = block_guard.response_sender().unbounded_send(event);
-		};
+				let _ = block_guard.response_sender().unbounded_send(event);
+			};
+			executor.spawn_blocking("substrate-rpc-subscription", Some("rpc"), fut.boxed());

-		self.executor.spawn("substrate-rpc-subscription", Some("rpc"), fut.boxed());
+			rp
+		});

-		rp
+		result
+			.await
+			.unwrap_or_else(|_| ResponsePayload::success(MethodResponse::LimitReached))
 	}

-	fn chain_head_unstable_header(
+	async fn chain_head_unstable_header(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		hash: Block::Hash,
 	) -> Result<Option<String>, ChainHeadRpcError> {
-		let _block_guard = match self.subscriptions.lock_block(&follow_subscription, hash, 1) {
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			return Ok(None);
+		}
+
+		let block_guard = match self.subscriptions.lock_block(&follow_subscription, hash, 1) {
 			Ok(block) => block,
 			Err(SubscriptionManagementError::SubscriptionAbsent) |
 			Err(SubscriptionManagementError::ExceededLimits) => return Ok(None),
@@ -296,19 +340,35 @@ where
 			Err(_) => return Err(ChainHeadRpcError::InvalidBlock.into()),
 		};

-		self.client
-			.header(hash)
-			.map(|opt_header| opt_header.map(|h| hex_string(&h.encode())))
-			.map_err(|err| ChainHeadRpcError::InternalError(err.to_string()))
+		let client = self.client.clone();
+		let result = spawn_blocking(&self.executor, async move {
+			let _block_guard = block_guard;
+
+			client
+				.header(hash)
+				.map(|opt_header| opt_header.map(|h| hex_string(&h.encode())))
+				.map_err(|err| ChainHeadRpcError::InternalError(err.to_string()))
+		});
+		result.await.unwrap_or_else(|_| Ok(None))
 	}

-	fn chain_head_unstable_storage(
+	async fn chain_head_unstable_storage(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		hash: Block::Hash,
 		items: Vec<StorageQuery<String>>,
 		child_trie: Option<String>,
 	) -> ResponsePayload<'static, MethodResponse> {
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			// The spec says to return `LimitReached` if the follow subscription is invalid or
+			// stale.
+			return ResponsePayload::success(MethodResponse::LimitReached);
+		}
+
 		// Gain control over parameter parsing and returned error.
 		let items = match items
 			.into_iter()
@@ -357,25 +417,25 @@ where
 		let mut items = items;
 		items.truncate(num_operations);

-		let (rp, rp_is_success) = method_started_response(operation_id, Some(discarded));
-
+		let (rp, rp_fut) = method_started_response(operation_id, Some(discarded));
 		let fut = async move {
-			// Events should only by generated
-			// if the response was successfully propagated.
-			if rp_is_success.await.is_err() {
+			// Wait for the server to send out the response and if it produces an error no event
+			// should be generated.
+			if rp_fut.await.is_err() {
 				return;
 			}
+
 			storage_client.generate_events(block_guard, hash, items, child_trie).await;
 		};
-
 		self.executor
 			.spawn_blocking("substrate-rpc-subscription", Some("rpc"), fut.boxed());

 		rp
 	}

-	fn chain_head_unstable_call(
+	async fn chain_head_unstable_call(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		hash: Block::Hash,
 		function: String,
@@ -386,6 +446,15 @@ where
 			Err(err) => return ResponsePayload::error(err),
 		};

+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			// The spec says to return `LimitReached` if the follow subscription is invalid or
+			// stale.
+			return ResponsePayload::success(MethodResponse::LimitReached);
+		}
+
 		let mut block_guard = match self.subscriptions.lock_block(&follow_subscription, hash, 1) {
 			Ok(block) => block,
 			Err(SubscriptionManagementError::SubscriptionAbsent) |
@@ -408,44 +477,53 @@ where
 		}

 		let operation_id = block_guard.operation().operation_id();
-		let event = self
-			.client
-			.executor()
-			.call(hash, &function, &call_parameters, CallContext::Offchain)
-			.map(|result| {
-				FollowEvent::<Block::Hash>::OperationCallDone(OperationCallDone {
-					operation_id: operation_id.clone(),
-					output: hex_string(&result),
-				})
-			})
-			.unwrap_or_else(|error| {
-				FollowEvent::<Block::Hash>::OperationError(OperationError {
-					operation_id: operation_id.clone(),
-					error: error.to_string(),
-				})
-			});
-
-		let (rp, rp_fut) = method_started_response(operation_id, None);
+		let client = self.client.clone();

+		let (rp, rp_fut) = method_started_response(operation_id.clone(), None);
 		let fut = async move {
-			// Events should only by generated
-			// if the response was successfully propagated.
+			// Wait for the server to send out the response and if it produces an error no event
+			// should be generated.
 			if rp_fut.await.is_err() {
-				return;
+				return
 			}
+
+			let event = client
+				.executor()
+				.call(hash, &function, &call_parameters, CallContext::Offchain)
+				.map(|result| {
+					FollowEvent::<Block::Hash>::OperationCallDone(OperationCallDone {
+						operation_id: operation_id.clone(),
+						output: hex_string(&result),
+					})
+				})
+				.unwrap_or_else(|error| {
+					FollowEvent::<Block::Hash>::OperationError(OperationError {
+						operation_id: operation_id.clone(),
+						error: error.to_string(),
+					})
+				});
+
 			let _ = block_guard.response_sender().unbounded_send(event);
 		};
-
-		self.executor.spawn("substrate-rpc-subscription", Some("rpc"), fut.boxed());
+		self.executor
+			.spawn_blocking("substrate-rpc-subscription", Some("rpc"), fut.boxed());

 		rp
 	}

-	fn chain_head_unstable_unpin(
+	async fn chain_head_unstable_unpin(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		hash_or_hashes: ListOrValue<Block::Hash>,
 	) -> Result<(), ChainHeadRpcError> {
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			return Ok(());
+		}
+
 		let result = match hash_or_hashes {
 			ListOrValue::Value(hash) =>
 				self.subscriptions.unpin_blocks(&follow_subscription, [hash]),
@@ -469,11 +547,19 @@ where
 		}
 	}

-	fn chain_head_unstable_continue(
+	async fn chain_head_unstable_continue(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		operation_id: String,
 	) -> Result<(), ChainHeadRpcError> {
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			return Ok(())
+		}
+
 		let Some(operation) = self.subscriptions.get_operation(&follow_subscription, &operation_id)
 		else {
 			return Ok(())
@@ -487,11 +573,19 @@ where
 		}
 	}

-	fn chain_head_unstable_stop_operation(
+	async fn chain_head_unstable_stop_operation(
 		&self,
+		connection_details: ConnectionDetails,
 		follow_subscription: String,
 		operation_id: String,
 	) -> Result<(), ChainHeadRpcError> {
+		if !self
+			.subscriptions
+			.contains_subscription(connection_details.id(), &follow_subscription)
+		{
+			return Ok(())
+		}
+
 		let Some(operation) = self.subscriptions.get_operation(&follow_subscription, &operation_id)
 		else {
 			return Ok(())
@@ -510,3 +604,26 @@ fn method_started_response(
 	let rp = MethodResponse::Started(MethodResponseStarted { operation_id, discarded_items });
 	ResponsePayload::success(rp).notify_on_completion()
 }
+
+/// Spawn a blocking future on the provided executor and return the result on a oneshot channel.
+///
+/// This is a wrapper to extract the result of a `executor.spawn_blocking` future.
+fn spawn_blocking<R>(
+	executor: &SubscriptionTaskExecutor,
+	fut: impl std::future::Future<Output = R> + Send + 'static,
+) -> oneshot::Receiver<R>
+where
+	R: Send + 'static,
+{
+	let (tx, rx) = oneshot::channel();
+
+	let blocking_fut = async move {
+		let result = fut.await;
+		// Send the result back on the channel.
+		let _ = tx.send(result);
+	};
+
+	executor.spawn_blocking("substrate-rpc-subscription", Some("rpc"), blocking_fut.boxed());
+
+	rx
+}