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_follow.rs

@@ -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,
@@ -50,6 +51,11 @@ use std::{
 /// The maximum number of finalized blocks provided by the
 /// `Initialized` event.
 const MAX_FINALIZED_BLOCKS: usize = 16;
+/// The size of the LRU cache for pruned blocks.
+///
+/// This is the exact value of the total number of pinned blocks, and ensures
+/// that all active pruned block hashes (if any) are kept in memory.
+const LRU_CACHE_SIZE: u32 = 512;
 
 use super::subscription::InsertedSubscriptionData;
 
@@ -67,6 +73,8 @@ pub struct ChainHeadFollower<BE: Backend<Block>, Block: BlockT, Client> {
 	sub_id: String,
 	/// The best reported block by this subscription.
 	best_block_cache: 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 +86,15 @@ 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,
+			best_block_cache: None,
+			pruned_blocks: LruMap::new(ByLength::new(LRU_CACHE_SIZE)),
+		}
 	}
 }
 
@@ -434,10 +450,9 @@ where
 	///
 	/// 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 +462,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 +482,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,7 +502,7 @@ 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,
@@ -495,15 +511,64 @@ where
 
 		match self.best_block_cache {
 			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 `NewBlock` 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.best_block_cache = 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 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 reported as pruned. Therefore, we need to signal a new
-				// best block event before submitting the finalized event.
+				// 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 +580,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.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)
 			},
 			None => {
 				events.push(finalized_event);
@@ -539,7 +604,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 +620,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 +706,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;
 	}
 }