chainHead: Report unique hashes for pruned blocks

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

@@ -42,6 +42,7 @@ array-bytes = "6.1"
 log = { workspace = true, default-features = true }
 futures-util = { version = "0.3.30", default-features = false }
 rand = "0.8.5"
+schnellru = "0.2.1"
 
 [dev-dependencies]
 serde_json = { workspace = true, default-features = true }

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

@@ -70,7 +70,7 @@ pub struct ChainHeadConfig {
 /// Maximum pinned blocks across all connections.
 /// This number is large enough to consider immediate blocks.
 /// Note: This should never exceed the `PINNING_CACHE_SIZE` from client/db.
-const MAX_PINNED_BLOCKS: usize = 512;
+pub(crate) const MAX_PINNED_BLOCKS: usize = 512;
 
 /// Any block of any subscription should not be pinned more than
 /// this constant. When a subscription contains a block older than this,

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

@@ -19,7 +19,7 @@
 //! Implementation of the `chainHead_follow` method.
 
 use crate::chain_head::{
-	chain_head::LOG_TARGET,
+	chain_head::{LOG_TARGET, MAX_PINNED_BLOCKS},
 	event::{
 		BestBlockChanged, Finalized, FollowEvent, Initialized, NewBlock, RuntimeEvent,
 		RuntimeVersionEvent,
@@ -37,6 +37,7 @@ use sc_client_api::{
 	Backend, BlockBackend, BlockImportNotification, BlockchainEvents, FinalityNotification,
 };
 use sc_rpc::utils::to_sub_message;
+use schnellru::{ByLength, LruMap};
 use sp_api::CallApiAt;
 use sp_blockchain::{
 	Backend as BlockChainBackend, Error as BlockChainError, HeaderBackend, HeaderMetadata, Info,
@@ -66,7 +67,9 @@ pub struct ChainHeadFollower<BE: Backend<Block>, Block: BlockT, Client> {
 	/// Subscription ID.
 	sub_id: String,
 	/// The best reported block by this subscription.
-	best_block_cache: Option<Block::Hash>,
+	current_best_block: Option<Block::Hash>,
+	/// LRU cache of pruned blocks.
+	pruned_blocks: LruMap<Block::Hash, ()>,
 }
 
 impl<BE: Backend<Block>, Block: BlockT, Client> ChainHeadFollower<BE, Block, Client> {
@@ -78,7 +81,17 @@ impl<BE: Backend<Block>, Block: BlockT, Client> ChainHeadFollower<BE, Block, Cli
 		with_runtime: bool,
 		sub_id: String,
 	) -> Self {
-		Self { client, backend, sub_handle, with_runtime, sub_id, best_block_cache: None }
+		Self {
+			client,
+			backend,
+			sub_handle,
+			with_runtime,
+			sub_id,
+			current_best_block: None,
+			pruned_blocks: LruMap::new(ByLength::new(
+				MAX_PINNED_BLOCKS.try_into().unwrap_or(u32::MAX),
+			)),
+		}
 	}
 }
 
@@ -295,7 +308,7 @@ where
 		let best_block_hash = startup_point.best_hash;
 		if best_block_hash != finalized_block_hash {
 			let best_block = FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash });
-			self.best_block_cache = Some(best_block_hash);
+			self.current_best_block = Some(best_block_hash);
 			finalized_block_descendants.push(best_block);
 		};
 
@@ -327,19 +340,19 @@ where
 		let best_block_event =
 			FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash: block_hash });
 
-		match self.best_block_cache {
+		match self.current_best_block {
 			Some(block_cache) => {
 				// The RPC layer has not reported this block as best before.
 				// Note: This handles the race with the finalized branch.
 				if block_cache != block_hash {
-					self.best_block_cache = Some(block_hash);
+					self.current_best_block = Some(block_hash);
 					vec![new_block, best_block_event]
 				} else {
 					vec![new_block]
 				}
 			},
 			None => {
-				self.best_block_cache = Some(block_hash);
+				self.current_best_block = Some(block_hash);
 				vec![new_block, best_block_event]
 			},
 		}
@@ -408,7 +421,7 @@ where
 				// When the node falls out of sync and then syncs up to the tip of the chain, it can
 				// happen that we skip notifications. Then it is better to terminate the connection
 				// instead of trying to send notifications for all missed blocks.
-				if let Some(best_block_hash) = self.best_block_cache {
+				if let Some(best_block_hash) = self.current_best_block {
 					let ancestor = sp_blockchain::lowest_common_ancestor(
 						&*self.client,
 						*hash,
@@ -431,13 +444,10 @@ where
 	}
 
 	/// Get all pruned block hashes from the provided stale heads.
-	///
-	/// The result does not include hashes from `to_ignore`.
 	fn get_pruned_hashes(
-		&self,
+		&mut self,
 		stale_heads: &[Block::Hash],
 		last_finalized: Block::Hash,
-		to_ignore: &mut HashSet<Block::Hash>,
 	) -> Result<Vec<Block::Hash>, SubscriptionManagementError> {
 		let blockchain = self.backend.blockchain();
 		let mut pruned = Vec::new();
@@ -447,11 +457,13 @@ where
 
 			// Collect only blocks that are not part of the canonical chain.
 			pruned.extend(tree_route.enacted().iter().filter_map(|block| {
-				if !to_ignore.remove(&block.hash) {
-					Some(block.hash)
-				} else {
-					None
+				if self.pruned_blocks.get(&block.hash).is_some() {
+					// The block was already reported as pruned.
+					return None
 				}
+
+				self.pruned_blocks.insert(block.hash, ());
+				Some(block.hash)
 			}))
 		}
 
@@ -465,7 +477,6 @@ where
 	fn handle_finalized_blocks(
 		&mut self,
 		notification: FinalityNotification<Block>,
-		to_ignore: &mut HashSet<Block::Hash>,
 		startup_point: &StartupPoint<Block>,
 	) -> Result<Vec<FollowEvent<Block::Hash>>, SubscriptionManagementError> {
 		let last_finalized = notification.hash;
@@ -486,24 +497,73 @@ where
 		// Report all pruned blocks from the notification that are not
 		// part of the fork we need to ignore.
 		let pruned_block_hashes =
-			self.get_pruned_hashes(&notification.stale_heads, last_finalized, to_ignore)?;
+			self.get_pruned_hashes(&notification.stale_heads, last_finalized)?;
 
 		let finalized_event = FollowEvent::Finalized(Finalized {
 			finalized_block_hashes,
 			pruned_block_hashes: pruned_block_hashes.clone(),
 		});
 
-		match self.best_block_cache {
+		match self.current_best_block {
 			Some(block_cache) => {
-				// If the best block wasn't pruned, we are done here.
-				if !pruned_block_hashes.iter().any(|hash| *hash == block_cache) {
+				// We need to generate a `BestBlock` event for the finalized block when:
+				// - (i) the last reported best block was pruned
+				// - (ii) the last reported best block is on a fork that will be pruned in the
+				//   future.
+				// Note: pruning happens on level n - 1.
+
+				// Best block already generated.
+				if block_cache == last_finalized {
 					events.push(finalized_event);
+					return Ok(events);
+				}
+
+				// Checking if the block was pruned is faster than computing the route tree.
+				let was_pruned = pruned_block_hashes.iter().any(|hash| *hash == block_cache);
+				if was_pruned {
+					// We need to generate a best block event.
+					let best_block_hash = self.client.info().best_hash;
+
+					// Defensive check against state missmatch.
+					if best_block_hash == block_cache {
+						// The client doest not have any new information about the best block.
+						// The information from `.info()` is updated from the DB as the last
+						// step of the finalization and it should be up to date.
+						// If the info is outdated, there is nothing the RPC can do for now.
+						error!(
+							target: LOG_TARGET,
+							"[follow][id={:?}] Client does not contain different best block",
+							self.sub_id,
+						);
+						events.push(finalized_event);
+						return Ok(events);
+					}
+
+					// The RPC needs to also submit a new best block changed before the
+					// finalized event.
+					self.current_best_block = Some(best_block_hash);
+					let best_block_event =
+						FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash });
+					events.extend([best_block_event, finalized_event]);
 					return Ok(events)
 				}
 
-				// The best block is reported as pruned. Therefore, we need to signal a new
-				// best block event before submitting the finalized event.
+				// The best block was not pruned, however it might be on a fork that will be pruned.
+				let tree_route = sp_blockchain::tree_route(
+					self.backend.blockchain(),
+					last_finalized,
+					block_cache,
+				)?;
+
+				// The best block is a descendent of the finalized block.
+				if tree_route.retracted().is_empty() {
+					events.push(finalized_event);
+					return Ok(events)
+				}
+
+				// The best block is on a fork that will be pruned.
 				let best_block_hash = self.client.info().best_hash;
+				// Defensive check against state missmatch.
 				if best_block_hash == block_cache {
 					// The client doest not have any new information about the best block.
 					// The information from `.info()` is updated from the DB as the last
@@ -515,16 +575,16 @@ where
 						self.sub_id,
 					);
 					events.push(finalized_event);
-					Ok(events)
-				} else {
-					// The RPC needs to also submit a new best block changed before the
-					// finalized event.
-					self.best_block_cache = Some(best_block_hash);
-					let best_block_event =
-						FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash });
-					events.extend([best_block_event, finalized_event]);
-					Ok(events)
+					return Ok(events);
 				}
+
+				// The RPC needs to also submit a new best block changed before the
+				// finalized event.
+				self.current_best_block = Some(best_block_hash);
+				let best_block_event =
+					FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash });
+				events.extend([best_block_event, finalized_event]);
+				Ok(events)
 			},
 			None => {
 				events.push(finalized_event);
@@ -539,7 +599,6 @@ where
 		&mut self,
 		startup_point: &StartupPoint<Block>,
 		mut stream: EventStream,
-		mut to_ignore: HashSet<Block::Hash>,
 		sink: SubscriptionSink,
 		rx_stop: oneshot::Receiver<()>,
 	) where
@@ -556,7 +615,7 @@ where
 				NotificationType::NewBlock(notification) =>
 					self.handle_import_blocks(notification, &startup_point),
 				NotificationType::Finalized(notification) =>
-					self.handle_finalized_blocks(notification, &mut to_ignore, &startup_point),
+					self.handle_finalized_blocks(notification, &startup_point),
 				NotificationType::MethodResponse(notification) => Ok(vec![notification]),
 			};
 
@@ -642,7 +701,10 @@ where
 		let merged = tokio_stream::StreamExt::merge(merged, stream_responses);
 		let stream = stream::once(futures::future::ready(initial)).chain(merged);
 
-		self.submit_events(&startup_point, stream.boxed(), pruned_forks, sink, sub_data.rx_stop)
-			.await;
+		// These are the pruned blocks that we should not report again.
+		for pruned in pruned_forks {
+			self.pruned_blocks.insert(pruned, ());
+		}
+		self.submit_events(&startup_point, stream.boxed(), sink, sub_data.rx_stop).await;
 	}
 }