From d9953487ab15118d3e76ea99de2cbe90af9d3440 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 15:08:29 +0200
Subject: [PATCH 01/16] chainHead: Report unique pruned hashes

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../client/rpc-spec-v2/src/chain_head/chain_head_follow.rs     | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index afa99f3aa164..4b9a32650bb8 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -441,13 +441,14 @@ where
 	) -> Result<Vec<Block::Hash>, SubscriptionManagementError> {
 		let blockchain = self.backend.blockchain();
 		let mut pruned = Vec::new();
+		let mut unique_hashes = HashSet::new();
 
 		for stale_head in stale_heads {
 			let tree_route = sp_blockchain::tree_route(blockchain, last_finalized, *stale_head)?;
 
 			// 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) {
+				if !to_ignore.remove(&block.hash) && unique_hashes.insert(block.hash) {
 					Some(block.hash)
 				} else {
 					None

From 00fd806ed73f545894c7bb235011397407f4d11d Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 16:16:37 +0200
Subject: [PATCH 02/16] chainHead: Test unique pruned blocks

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 262 ++++++++++++++++++
 1 file changed, 262 insertions(+)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index 4d9dfb24e0a9..4f5743256094 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3297,3 +3297,265 @@ async fn storage_closest_merkle_value() {
 		merkle_values_rhs.get(&hex_string(b":AAAA")).unwrap()
 	);
 }
+
+#[tokio::test]
+async fn follow_unique_pruned_blocks() {
+	let builder = TestClientBuilder::new();
+	let backend = builder.backend();
+	let mut client = Arc::new(builder.build());
+
+	let api = ChainHead::new(
+		client.clone(),
+		backend,
+		Arc::new(TaskExecutor::default()),
+		ChainHeadConfig {
+			global_max_pinned_blocks: MAX_PINNED_BLOCKS,
+			subscription_max_pinned_duration: Duration::from_secs(MAX_PINNED_SECS),
+			subscription_max_ongoing_operations: MAX_OPERATIONS,
+			operation_max_storage_items: MAX_PAGINATION_LIMIT,
+		},
+	)
+	.into_rpc();
+
+	let finalized_hash = client.info().finalized_hash;
+	let mut sub = api.subscribe_unbounded("chainHead_unstable_follow", [false]).await.unwrap();
+
+	// Initialized must always be reported first.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::Initialized(Initialized {
+		finalized_block_hashes: vec![format!("{:?}", finalized_hash)],
+		finalized_block_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+
+	// Block tree:
+	//
+	// finalized -> block 1 -> block 2 -> block 3
+	//
+	//                      -> block 2 -> block 4
+	//
+	//           -> block 1 -> block 2_f -> block 5
+	//                                    ^^^ finalized
+	//
+	// The block 4 is needed on the longest chain because we want the
+	// best block 2 to be reported as pruned. Pruning is happening at
+	// height (N - 1), where N is the finalized block number.
+
+	let block_1 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(client.chain_info().genesis_hash)
+		.with_parent_block_number(0)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_1_hash = block_1.hash();
+	println!("block_1_hash: {:?}", block_1_hash);
+	client.import(BlockOrigin::Own, block_1.clone()).await.unwrap();
+
+	let block_2_f = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_1_hash)
+		.with_parent_block_number(1)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_2_f_hash = block_2_f.hash();
+	println!("block_2_f_hash: {:?}", block_2_f_hash);
+	client.import(BlockOrigin::Own, block_2_f.clone()).await.unwrap();
+
+	let block_5 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_2_f_hash)
+		.with_parent_block_number(2)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_5_hash = block_5.hash();
+	println!("block_5_hash: {:?}", block_5_hash);
+
+	client.import(BlockOrigin::Own, block_5.clone()).await.unwrap();
+
+	// Import block 2 as best on the fork.
+	let mut block_builder = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_1_hash)
+		.with_parent_block_number(1)
+		.build()
+		.unwrap();
+	// This push is required as otherwise block 3 has the same hash as block 2 and won't get
+	// imported
+	block_builder
+		.push_transfer(Transfer {
+			from: AccountKeyring::Alice.into(),
+			to: AccountKeyring::Ferdie.into(),
+			amount: 41,
+			nonce: 0,
+		})
+		.unwrap();
+	let block_2 = block_builder.build().unwrap().block;
+	let block_2_hash = block_2.header.hash();
+	println!("block_2_hash: {:?}", block_2_hash);
+
+	client.import_as_best(BlockOrigin::Own, block_2.clone()).await.unwrap();
+
+	let block_3 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_2_hash)
+		.with_parent_block_number(2)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_3_hash = block_3.hash();
+	println!("block_3_hash: {:?}", block_3_hash);
+
+	client.import(BlockOrigin::Own, block_3.clone()).await.unwrap();
+
+	// Fork block 4.
+	let mut block_builder = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_2_hash)
+		.with_parent_block_number(2)
+		.build()
+		.unwrap();
+	// This push is required as otherwise block 3 has the same hash as block 2 and won't get
+	// imported
+	block_builder
+		.push_transfer(Transfer {
+			from: AccountKeyring::Alice.into(),
+			to: AccountKeyring::Ferdie.into(),
+			amount: 41,
+			nonce: 1,
+		})
+		.unwrap();
+	let block_4 = block_builder.build().unwrap().block;
+	let block_4_hash = block_4.header.hash();
+	println!("block_4_hash: {:?}", block_4_hash);
+
+	client.import_as_best(BlockOrigin::Own, block_4.clone()).await.unwrap();
+
+	// Check block 1.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_1_hash),
+		parent_block_hash: format!("{:?}", finalized_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_1_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Check block 2f.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_2_f_hash),
+		parent_block_hash: format!("{:?}", block_1_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_2_f_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Check block 5.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_5_hash),
+		parent_block_hash: format!("{:?}", block_2_f_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_5_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Check block 2, that we imported as custom best.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_2_hash),
+		parent_block_hash: format!("{:?}", block_1_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_2_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Check block 3.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_3_hash),
+		parent_block_hash: format!("{:?}", block_2_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_3_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Check block 4.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_4_hash),
+		parent_block_hash: format!("{:?}", block_2_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_4_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Finalize the block 5 from the fork.
+	client.finalize_block(block_5_hash, None).unwrap();
+
+	// Conclusion:
+	// - Block 5 is not reported as bewst; block 4 is the last best block.
+	//
+	// - block 2 is not pruned (it should be reported as pruned at N - 1, so its ok)
+
+	// Expect to report the best block changed before the finalized event.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_5_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Block 2 must be reported as pruned, even if it was the previous best.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::Finalized(Finalized {
+		finalized_block_hashes: vec![
+			format!("{:?}", block_1_hash),
+			format!("{:?}", block_2_f_hash),
+			format!("{:?}", block_5_hash),
+		],
+		pruned_block_hashes: vec![format!("{:?}", block_2_hash)],
+	});
+	assert_eq!(event, expected);
+
+	// Pruned hash can be unpinned.
+	let sub_id = sub.subscription_id();
+	let sub_id = serde_json::to_string(&sub_id).unwrap();
+	let hash = format!("{:?}", block_2_hash);
+	let _res: () = api.call("chainHead_unstable_unpin", rpc_params![&sub_id, &hash]).await.unwrap();
+}

From 494a1eaa256d370a4640a01ab222bb61fb4f63d9 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 16:43:08 +0200
Subject: [PATCH 03/16] chainHead: Generate best block event even for unpruned
 forks

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../src/chain_head/chain_head_follow.rs       | 75 +++++++++++++++----
 1 file changed, 62 insertions(+), 13 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 4b9a32650bb8..bc9bfe69f77c 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -496,15 +496,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
@@ -516,16 +565,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);

From 10b143a5e99393df23677306f88538e7b6fe0448 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 16:57:20 +0200
Subject: [PATCH 04/16] chainHead/tests: Add one more block to trigger pruning

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 47 ++++++++++++++++++-
 1 file changed, 46 insertions(+), 1 deletion(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index 4f5743256094..a0e663868624 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3549,7 +3549,7 @@ async fn follow_unique_pruned_blocks() {
 			format!("{:?}", block_2_f_hash),
 			format!("{:?}", block_5_hash),
 		],
-		pruned_block_hashes: vec![format!("{:?}", block_2_hash)],
+		pruned_block_hashes: vec![],
 	});
 	assert_eq!(event, expected);
 
@@ -3558,4 +3558,49 @@ async fn follow_unique_pruned_blocks() {
 	let sub_id = serde_json::to_string(&sub_id).unwrap();
 	let hash = format!("{:?}", block_2_hash);
 	let _res: () = api.call("chainHead_unstable_unpin", rpc_params![&sub_id, &hash]).await.unwrap();
+
+	// Check block 6.
+	let block_6 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_5_hash)
+		.with_parent_block_number(3)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_6_hash = block_6.hash();
+	println!("block_6_hash: {:?}", block_6_hash);
+	client.import(BlockOrigin::Own, block_6.clone()).await.unwrap();
+
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_6_hash),
+		parent_block_hash: format!("{:?}", block_5_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_6_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Finalize the block 6.
+	client.finalize_block(block_6_hash, None).unwrap();
+	println!("Finalized block 6");
+
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::Finalized(Finalized {
+		finalized_block_hashes: vec![format!("{:?}", block_6_hash)],
+		pruned_block_hashes: vec![
+			format!("{:?}", block_2_hash),
+			format!("{:?}", block_3_hash),
+			format!("{:?}", block_4_hash),
+		],
+	});
+	assert_eq!(event, expected);
 }

From c840c8bb8a9d5926c75d3d5e60303747c6524186 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 17:05:40 +0200
Subject: [PATCH 05/16] chainHead/tests: An extra block to validate pruned
 blocks are reported

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 129 +++++++++++++-----
 1 file changed, 95 insertions(+), 34 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index a0e663868624..076551ea2d8e 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3333,9 +3333,9 @@ async fn follow_unique_pruned_blocks() {
 	//
 	// finalized -> block 1 -> block 2 -> block 3
 	//
-	//                      -> block 2 -> block 4
+	//                      -> block 2 -> block 4 -> block 5
 	//
-	//           -> block 1 -> block 2_f -> block 5
+	//           -> block 1 -> block 2_f -> block 6     -> (added later) block 7 -> block 8
 	//                                    ^^^ finalized
 	//
 	// The block 4 is needed on the longest chain because we want the
@@ -3366,7 +3366,7 @@ async fn follow_unique_pruned_blocks() {
 	println!("block_2_f_hash: {:?}", block_2_f_hash);
 	client.import(BlockOrigin::Own, block_2_f.clone()).await.unwrap();
 
-	let block_5 = BlockBuilderBuilder::new(&*client)
+	let block_6 = BlockBuilderBuilder::new(&*client)
 		.on_parent_block(block_2_f_hash)
 		.with_parent_block_number(2)
 		.build()
@@ -3374,10 +3374,10 @@ async fn follow_unique_pruned_blocks() {
 		.build()
 		.unwrap()
 		.block;
-	let block_5_hash = block_5.hash();
-	println!("block_5_hash: {:?}", block_5_hash);
+	let block_6_hash = block_6.hash();
+	println!("block_6: {:?}", block_6_hash);
 
-	client.import(BlockOrigin::Own, block_5.clone()).await.unwrap();
+	client.import(BlockOrigin::Own, block_6.clone()).await.unwrap();
 
 	// Import block 2 as best on the fork.
 	let mut block_builder = BlockBuilderBuilder::new(&*client)
@@ -3398,7 +3398,6 @@ async fn follow_unique_pruned_blocks() {
 	let block_2 = block_builder.build().unwrap().block;
 	let block_2_hash = block_2.header.hash();
 	println!("block_2_hash: {:?}", block_2_hash);
-
 	client.import_as_best(BlockOrigin::Own, block_2.clone()).await.unwrap();
 
 	let block_3 = BlockBuilderBuilder::new(&*client)
@@ -3411,7 +3410,6 @@ async fn follow_unique_pruned_blocks() {
 		.block;
 	let block_3_hash = block_3.hash();
 	println!("block_3_hash: {:?}", block_3_hash);
-
 	client.import(BlockOrigin::Own, block_3.clone()).await.unwrap();
 
 	// Fork block 4.
@@ -3433,9 +3431,20 @@ async fn follow_unique_pruned_blocks() {
 	let block_4 = block_builder.build().unwrap().block;
 	let block_4_hash = block_4.header.hash();
 	println!("block_4_hash: {:?}", block_4_hash);
-
 	client.import_as_best(BlockOrigin::Own, block_4.clone()).await.unwrap();
 
+	let block_5 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_4_hash)
+		.with_parent_block_number(3)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_5_hash = block_5.hash();
+	println!("block_5_hash: {:?}", block_5_hash);
+	client.import(BlockOrigin::Own, block_5.clone()).await.unwrap();
+
 	// Check block 1.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::NewBlock(NewBlock {
@@ -3466,10 +3475,10 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 
-	// Check block 5.
+	// Check block 6.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_5_hash),
+		block_hash: format!("{:?}", block_6_hash),
 		parent_block_hash: format!("{:?}", block_2_f_hash),
 		new_runtime: None,
 		with_runtime: false,
@@ -3477,7 +3486,7 @@ async fn follow_unique_pruned_blocks() {
 	assert_eq!(event, expected);
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_5_hash),
+		best_block_hash: format!("{:?}", block_6_hash),
 	});
 	assert_eq!(event, expected);
 
@@ -3526,18 +3535,28 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 
-	// Finalize the block 5 from the fork.
-	client.finalize_block(block_5_hash, None).unwrap();
+	// Check block 5.
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_5_hash),
+		parent_block_hash: format!("{:?}", block_4_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_5_hash),
+	});
+	assert_eq!(event, expected);
 
-	// Conclusion:
-	// - Block 5 is not reported as bewst; block 4 is the last best block.
-	//
-	// - block 2 is not pruned (it should be reported as pruned at N - 1, so its ok)
+	// Finalize the block 6 from the fork.
+	client.finalize_block(block_6_hash, None).unwrap();
 
 	// Expect to report the best block changed before the finalized event.
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_5_hash),
+		best_block_hash: format!("{:?}", block_6_hash),
 	});
 	assert_eq!(event, expected);
 
@@ -3547,7 +3566,7 @@ async fn follow_unique_pruned_blocks() {
 		finalized_block_hashes: vec![
 			format!("{:?}", block_1_hash),
 			format!("{:?}", block_2_f_hash),
-			format!("{:?}", block_5_hash),
+			format!("{:?}", block_6_hash),
 		],
 		pruned_block_hashes: vec![],
 	});
@@ -3559,24 +3578,24 @@ async fn follow_unique_pruned_blocks() {
 	let hash = format!("{:?}", block_2_hash);
 	let _res: () = api.call("chainHead_unstable_unpin", rpc_params![&sub_id, &hash]).await.unwrap();
 
-	// Check block 6.
-	let block_6 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_5_hash)
+	// Check block 7.
+	let block_7 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_6_hash)
 		.with_parent_block_number(3)
 		.build()
 		.unwrap()
 		.build()
 		.unwrap()
 		.block;
-	let block_6_hash = block_6.hash();
-	println!("block_6_hash: {:?}", block_6_hash);
-	client.import(BlockOrigin::Own, block_6.clone()).await.unwrap();
+	let block_7_hash = block_7.hash();
+	println!("block_7: {:?}", block_7_hash);
+	client.import(BlockOrigin::Own, block_7.clone()).await.unwrap();
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	println!("event: {:?}", event);
 	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_6_hash),
-		parent_block_hash: format!("{:?}", block_5_hash),
+		block_hash: format!("{:?}", block_7_hash),
+		parent_block_hash: format!("{:?}", block_6_hash),
 		new_runtime: None,
 		with_runtime: false,
 	});
@@ -3584,22 +3603,64 @@ async fn follow_unique_pruned_blocks() {
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	println!("event: {:?}", event);
 	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_6_hash),
+		best_block_hash: format!("{:?}", block_7_hash),
 	});
 	assert_eq!(event, expected);
 
-	// Finalize the block 6.
-	client.finalize_block(block_6_hash, None).unwrap();
-	println!("Finalized block 6");
+	// Finalize the block 7.
+	client.finalize_block(block_7_hash, None).unwrap();
+	println!("Finalized block 7");
+
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::Finalized(Finalized {
+		finalized_block_hashes: vec![format!("{:?}", block_7_hash)],
+		pruned_block_hashes: vec![format!("{:?}", block_2_hash), format!("{:?}", block_3_hash)],
+	});
+	assert_eq!(event, expected);
+
+	// Check block 8.
+	let block_8 = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(block_7_hash)
+		.with_parent_block_number(4)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	let block_8_hash = block_8.hash();
+	println!("block_8_hash: {:?}", block_8_hash);
+	client.import(BlockOrigin::Own, block_8.clone()).await.unwrap();
+
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::NewBlock(NewBlock {
+		block_hash: format!("{:?}", block_8_hash),
+		parent_block_hash: format!("{:?}", block_7_hash),
+		new_runtime: None,
+		with_runtime: false,
+	});
+	assert_eq!(event, expected);
+	let event: FollowEvent<String> = get_next_event(&mut sub).await;
+	println!("event: {:?}", event);
+	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+		best_block_hash: format!("{:?}", block_8_hash),
+	});
+	assert_eq!(event, expected);
+
+	// Finalize the block 8.
+	client.finalize_block(block_8_hash, None).unwrap();
+	println!("Finalized block 8");
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
 	println!("event: {:?}", event);
 	let expected = FollowEvent::Finalized(Finalized {
-		finalized_block_hashes: vec![format!("{:?}", block_6_hash)],
+		finalized_block_hashes: vec![format!("{:?}", block_8_hash)],
 		pruned_block_hashes: vec![
 			format!("{:?}", block_2_hash),
-			format!("{:?}", block_3_hash),
+			// However only block 4 and 5 should be reported.
 			format!("{:?}", block_4_hash),
+			format!("{:?}", block_5_hash),
 		],
 	});
 	assert_eq!(event, expected);

From 3f3d8767d76a880051cf175123e2edc86833a014 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 18:23:29 +0200
Subject: [PATCH 06/16] chainHead: Use LRU for caching pruned blocks

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 Cargo.lock                                    |  1 +
 substrate/client/rpc-spec-v2/Cargo.toml       |  1 +
 .../src/chain_head/chain_head_follow.rs       | 45 +++++++++++++------
 .../rpc-spec-v2/src/chain_head/tests.rs       |  7 +--
 4 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 37c5a5177c04..6d7df9ee627b 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -16701,6 +16701,7 @@ dependencies = [
  "sc-transaction-pool",
  "sc-transaction-pool-api",
  "sc-utils",
+ "schnellru",
  "serde",
  "serde_json",
  "sp-api",
diff --git a/substrate/client/rpc-spec-v2/Cargo.toml b/substrate/client/rpc-spec-v2/Cargo.toml
index c62b3e789d38..2c87af30284b 100644
--- a/substrate/client/rpc-spec-v2/Cargo.toml
+++ b/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 }
diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index bc9bfe69f77c..6746b51bcf64 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/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,25 +450,25 @@ 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();
-		let mut unique_hashes = HashSet::new();
 
 		for stale_head in stale_heads {
 			let tree_route = sp_blockchain::tree_route(blockchain, last_finalized, *stale_head)?;
 
 			// 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) && unique_hashes.insert(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)
 			}))
 		}
 
@@ -466,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;
@@ -487,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,
@@ -589,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
@@ -606,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]),
 			};
 
@@ -692,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;
 	}
 }
diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index 076551ea2d8e..d5cb95b6bdc1 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3656,12 +3656,7 @@ async fn follow_unique_pruned_blocks() {
 	println!("event: {:?}", event);
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_8_hash)],
-		pruned_block_hashes: vec![
-			format!("{:?}", block_2_hash),
-			// However only block 4 and 5 should be reported.
-			format!("{:?}", block_4_hash),
-			format!("{:?}", block_5_hash),
-		],
+		pruned_block_hashes: vec![format!("{:?}", block_4_hash), format!("{:?}", block_5_hash)],
 	});
 	assert_eq!(event, expected);
 }

From 936434c721efb46f9d9c1306d0ce077c9cefc1c2 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 12 Mar 2024 18:35:25 +0200
Subject: [PATCH 07/16] chainHead/tests: Remove debug logs and add comment
 about test

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 38 +++++++++----------
 1 file changed, 17 insertions(+), 21 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index d5cb95b6bdc1..a16b337eead0 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3335,12 +3335,25 @@ async fn follow_unique_pruned_blocks() {
 	//
 	//                      -> block 2 -> block 4 -> block 5
 	//
-	//           -> block 1 -> block 2_f -> block 6     -> (added later) block 7 -> block 8
+	//           -> block 1 -> block 2_f -> block 6
 	//                                    ^^^ finalized
+	//                                                 -> block 7
+	//                                                    ^^^ finalized
+	//                                                            -> block 8
+	//                                                               ^^^ finalized
+	// The chainHead will see block 5 as the best block. However, the
+	// client will finalize the block 6, which is on another fork.
 	//
-	// The block 4 is needed on the longest chain because we want the
-	// best block 2 to be reported as pruned. Pruning is happening at
-	// height (N - 1), where N is the finalized block number.
+	// When the block 6 is finalized, block 2 block 3 block 4 and block 5 are placed on an invalid
+	// fork. However, pruning of blocks happens on level N - 1.
+	// Therefore, no pruned blocks are reported yet.
+	//
+	// When the block 7 is finalized, block 3 is detected as stale. At this step, block 2 and 3
+	// are reported as pruned.
+	//
+	// When the block 8 is finalized, block 5 block 4 and block 2 are detected as stale. However,
+	// only blocks 5 and 4 are reported as pruned. This is because the block 2 was previously
+	// reported.
 
 	let block_1 = BlockBuilderBuilder::new(&*client)
 		.on_parent_block(client.chain_info().genesis_hash)
@@ -3351,7 +3364,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_1_hash = block_1.hash();
-	println!("block_1_hash: {:?}", block_1_hash);
 	client.import(BlockOrigin::Own, block_1.clone()).await.unwrap();
 
 	let block_2_f = BlockBuilderBuilder::new(&*client)
@@ -3363,7 +3375,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_2_f_hash = block_2_f.hash();
-	println!("block_2_f_hash: {:?}", block_2_f_hash);
 	client.import(BlockOrigin::Own, block_2_f.clone()).await.unwrap();
 
 	let block_6 = BlockBuilderBuilder::new(&*client)
@@ -3375,7 +3386,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_6_hash = block_6.hash();
-	println!("block_6: {:?}", block_6_hash);
 
 	client.import(BlockOrigin::Own, block_6.clone()).await.unwrap();
 
@@ -3397,7 +3407,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap();
 	let block_2 = block_builder.build().unwrap().block;
 	let block_2_hash = block_2.header.hash();
-	println!("block_2_hash: {:?}", block_2_hash);
 	client.import_as_best(BlockOrigin::Own, block_2.clone()).await.unwrap();
 
 	let block_3 = BlockBuilderBuilder::new(&*client)
@@ -3409,7 +3418,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_3_hash = block_3.hash();
-	println!("block_3_hash: {:?}", block_3_hash);
 	client.import(BlockOrigin::Own, block_3.clone()).await.unwrap();
 
 	// Fork block 4.
@@ -3430,7 +3438,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap();
 	let block_4 = block_builder.build().unwrap().block;
 	let block_4_hash = block_4.header.hash();
-	println!("block_4_hash: {:?}", block_4_hash);
 	client.import_as_best(BlockOrigin::Own, block_4.clone()).await.unwrap();
 
 	let block_5 = BlockBuilderBuilder::new(&*client)
@@ -3442,7 +3449,6 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_5_hash = block_5.hash();
-	println!("block_5_hash: {:?}", block_5_hash);
 	client.import(BlockOrigin::Own, block_5.clone()).await.unwrap();
 
 	// Check block 1.
@@ -3588,11 +3594,9 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_7_hash = block_7.hash();
-	println!("block_7: {:?}", block_7_hash);
 	client.import(BlockOrigin::Own, block_7.clone()).await.unwrap();
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::NewBlock(NewBlock {
 		block_hash: format!("{:?}", block_7_hash),
 		parent_block_hash: format!("{:?}", block_6_hash),
@@ -3601,7 +3605,6 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
 		best_block_hash: format!("{:?}", block_7_hash),
 	});
@@ -3609,10 +3612,8 @@ async fn follow_unique_pruned_blocks() {
 
 	// Finalize the block 7.
 	client.finalize_block(block_7_hash, None).unwrap();
-	println!("Finalized block 7");
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_7_hash)],
 		pruned_block_hashes: vec![format!("{:?}", block_2_hash), format!("{:?}", block_3_hash)],
@@ -3629,11 +3630,9 @@ async fn follow_unique_pruned_blocks() {
 		.unwrap()
 		.block;
 	let block_8_hash = block_8.hash();
-	println!("block_8_hash: {:?}", block_8_hash);
 	client.import(BlockOrigin::Own, block_8.clone()).await.unwrap();
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::NewBlock(NewBlock {
 		block_hash: format!("{:?}", block_8_hash),
 		parent_block_hash: format!("{:?}", block_7_hash),
@@ -3642,7 +3641,6 @@ async fn follow_unique_pruned_blocks() {
 	});
 	assert_eq!(event, expected);
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
 		best_block_hash: format!("{:?}", block_8_hash),
 	});
@@ -3650,10 +3648,8 @@ async fn follow_unique_pruned_blocks() {
 
 	// Finalize the block 8.
 	client.finalize_block(block_8_hash, None).unwrap();
-	println!("Finalized block 8");
 
 	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	println!("event: {:?}", event);
 	let expected = FollowEvent::Finalized(Finalized {
 		finalized_block_hashes: vec![format!("{:?}", block_8_hash)],
 		pruned_block_hashes: vec![format!("{:?}", block_4_hash), format!("{:?}", block_5_hash)],

From a33ed67e19e212d5e8f1d9252a1076de3b33fec4 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <60601340+lexnv@users.noreply.github.com>
Date: Fri, 15 Mar 2024 11:27:45 +0200
Subject: [PATCH 08/16] Update
 substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs

Co-authored-by: Sebastian Kunert <skunert49@gmail.com>
---
 .../client/rpc-spec-v2/src/chain_head/chain_head_follow.rs      | 2 --
 1 file changed, 2 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 6746b51bcf64..501724bc3fa4 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -447,8 +447,6 @@ where
 	}
 
 	/// Get all pruned block hashes from the provided stale heads.
-	///
-	/// The result does not include hashes from `to_ignore`.
 	fn get_pruned_hashes(
 		&mut self,
 		stale_heads: &[Block::Hash],

From 2c05bc58e8ca83745257d6baf5d098d31136f634 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <60601340+lexnv@users.noreply.github.com>
Date: Fri, 15 Mar 2024 13:34:59 +0200
Subject: [PATCH 09/16] Update
 substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs

---
 .../client/rpc-spec-v2/src/chain_head/chain_head_follow.rs      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 501724bc3fa4..8aa7bd126625 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -509,7 +509,7 @@ where
 
 		match self.best_block_cache {
 			Some(block_cache) => {
-				// We need to generate a `NewBlock` event for the finalized block when:
+				// 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.

From 384d22ca4c0fd8b53d4d7d069220252216d5693b Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Fri, 15 Mar 2024 15:40:42 +0200
Subject: [PATCH 10/16] chainHead: Rename best_block_cache to
 current_best_block

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../src/chain_head/chain_head_follow.rs       | 20 +++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 8aa7bd126625..83c0d669cb16 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -72,7 +72,7 @@ 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, ()>,
 }
@@ -92,7 +92,7 @@ impl<BE: Backend<Block>, Block: BlockT, Client> ChainHeadFollower<BE, Block, Cli
 			sub_handle,
 			with_runtime,
 			sub_id,
-			best_block_cache: None,
+			current_best_block: None,
 			pruned_blocks: LruMap::new(ByLength::new(LRU_CACHE_SIZE)),
 		}
 	}
@@ -311,7 +311,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);
 		};
 
@@ -343,19 +343,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]
 			},
 		}
@@ -424,7 +424,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,
@@ -507,7 +507,7 @@ where
 			pruned_block_hashes: pruned_block_hashes.clone(),
 		});
 
-		match self.best_block_cache {
+		match self.current_best_block {
 			Some(block_cache) => {
 				// We need to generate a `BestBlock` event for the finalized block when:
 				// - (i) the last reported best block was pruned
@@ -544,7 +544,7 @@ where
 
 					// The RPC needs to also submit a new best block changed before the
 					// finalized event.
-					self.best_block_cache = Some(best_block_hash);
+					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]);
@@ -583,7 +583,7 @@ where
 
 				// The RPC needs to also submit a new best block changed before the
 				// finalized event.
-				self.best_block_cache = Some(best_block_hash);
+				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]);

From 0b7ae73e30e6e16c54d403eae6c7363fe7976596 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Fri, 15 Mar 2024 15:43:39 +0200
Subject: [PATCH 11/16] chainHead: Use num pinned blocks for maximum LRU cache
 size

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../client/rpc-spec-v2/src/chain_head/chain_head.rs   |  2 +-
 .../rpc-spec-v2/src/chain_head/chain_head_follow.rs   | 11 ++++-------
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head.rs
index 2bda22b45239..aaac0b785315 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head.rs
+++ b/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,
diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 83c0d669cb16..9c6e33c63651 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/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,
@@ -51,11 +51,6 @@ 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;
 
@@ -93,7 +88,9 @@ impl<BE: Backend<Block>, Block: BlockT, Client> ChainHeadFollower<BE, Block, Cli
 			with_runtime,
 			sub_id,
 			current_best_block: None,
-			pruned_blocks: LruMap::new(ByLength::new(LRU_CACHE_SIZE)),
+			pruned_blocks: LruMap::new(ByLength::new(
+				MAX_PINNED_BLOCKS.try_into().unwrap_or(u32::MAX),
+			)),
 		}
 	}
 }

From e644f0aa76dce50193f0858a8de4c17522740704 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 19 Mar 2024 17:20:14 +0200
Subject: [PATCH 12/16] chainHead: Simplify new block generation logic

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../src/chain_head/chain_head_follow.rs       | 101 +++++-------------
 1 file changed, 27 insertions(+), 74 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index 9c6e33c63651..61eb90f58e9b 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -504,67 +504,25 @@ where
 			pruned_block_hashes: pruned_block_hashes.clone(),
 		});
 
-		match self.current_best_block {
-			Some(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 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.
+		if let Some(current_best_block) = self.current_best_block {
+			// The best reported block is in the pruned list. Report a new best block.
+			let is_in_pruned_list =
+				pruned_block_hashes.iter().any(|hash| *hash == current_best_block);
+			// The block is not the last finalized block.
+			//
+			// It can be either:
+			//  - a descendant of the last finalized block
+			//  - a block on a fork that will be pruned in the future.
+			//
+			// In those cases, we emit a new best block.
+			let is_not_last_finalized = current_best_block != last_finalized;
+
+			if is_in_pruned_list || is_not_last_finalized {
+				// 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 {
+				if best_block_hash == current_best_block {
 					// 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.
@@ -574,23 +532,18 @@ where
 						"[follow][id={:?}] Client does not contain different best block",
 						self.sub_id,
 					);
-					events.push(finalized_event);
-					return Ok(events);
+				} else {
+					// The RPC needs to also submit a new best block changed before the
+					// finalized event.
+					self.current_best_block = Some(best_block_hash);
+					events
+						.push(FollowEvent::BestBlockChanged(BestBlockChanged { best_block_hash }));
 				}
-
-				// 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);
-				Ok(events)
-			},
+			}
 		}
+
+		events.push(finalized_event);
+		Ok(events)
 	}
 
 	/// Submit the events from the provided stream to the RPC client

From 99bf9c90332f57d99333b64d9a5a49ef7003ea87 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 9 Apr 2024 15:04:22 +0300
Subject: [PATCH 13/16] chainHead/tests: Add max sub config for tests

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 substrate/client/rpc-spec-v2/src/chain_head/tests.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index b0cfb79a55fc..4766019ddf70 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3539,6 +3539,7 @@ async fn follow_unique_pruned_blocks() {
 			subscription_max_pinned_duration: Duration::from_secs(MAX_PINNED_SECS),
 			subscription_max_ongoing_operations: MAX_OPERATIONS,
 			operation_max_storage_items: MAX_PAGINATION_LIMIT,
+			max_follow_subscriptions_per_connection: MAX_FOLLOW_SUBSCRIPTIONS_PER_CONNECTION,
 		},
 	)
 	.into_rpc();

From c52d42d4419f8deaab5c0965a4a14886ac3626df Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 9 Apr 2024 16:57:15 +0300
Subject: [PATCH 14/16] Fix clippy

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 substrate/client/rpc-spec-v2/src/chain_head/tests.rs | 1 +
 1 file changed, 1 insertion(+)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index bdf14d7a3804..3347d6d8ac6d 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -3661,6 +3661,7 @@ async fn follow_unique_pruned_blocks() {
 			subscription_max_ongoing_operations: MAX_OPERATIONS,
 			operation_max_storage_items: MAX_PAGINATION_LIMIT,
 			max_follow_subscriptions_per_connection: MAX_FOLLOW_SUBSCRIPTIONS_PER_CONNECTION,
+			max_lagging_distance: MAX_LAGGING_DISTANCE,
 		},
 	)
 	.into_rpc();

From 0acd861a4cfa023dd704194f66f83af6068a4947 Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Mon, 15 Apr 2024 13:55:23 +0300
Subject: [PATCH 15/16] chainHead/follow: Simplify generate_init_events by
 storing the pruned hashes

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../src/chain_head/chain_head_follow.rs        | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
index db08455fc651..a753896b24c2 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/chain_head_follow.rs
@@ -309,18 +309,20 @@ where
 
 	/// Generate the initial events reported by the RPC `follow` method.
 	///
-	/// Returns the initial events that should be reported directly, together with pruned
-	/// block hashes that should be ignored by the `Finalized` event.
+	/// Returns the initial events that should be reported directly.
 	fn generate_init_events(
 		&mut self,
 		startup_point: &StartupPoint<Block>,
-	) -> Result<(Vec<FollowEvent<Block::Hash>>, HashSet<Block::Hash>), SubscriptionManagementError>
-	{
+	) -> Result<Vec<FollowEvent<Block::Hash>>, SubscriptionManagementError> {
 		let init = self.get_init_blocks_with_forks(startup_point.finalized_hash)?;
 
 		// The initialized event is the first one sent.
 		let initial_blocks = init.finalized_block_descendants;
 		let finalized_block_hashes = init.finalized_block_hashes;
+		// These are the pruned blocks that we should not report again.
+		for pruned in init.pruned_forks {
+			self.pruned_blocks.insert(pruned, ());
+		}
 
 		let finalized_block_hash = startup_point.finalized_hash;
 		let finalized_block_runtime = self.generate_runtime_event(finalized_block_hash, None);
@@ -355,7 +357,7 @@ where
 			finalized_block_descendants.push(best_block);
 		};
 
-		Ok((finalized_block_descendants, init.pruned_forks))
+		Ok(finalized_block_descendants)
 	}
 
 	/// Generate the "NewBlock" event and potentially the "BestBlockChanged" event for the
@@ -687,7 +689,7 @@ where
 			.map(|response| NotificationType::MethodResponse(response));
 
 		let startup_point = StartupPoint::from(self.client.info());
-		let (initial_events, pruned_forks) = match self.generate_init_events(&startup_point) {
+		let initial_events = match self.generate_init_events(&startup_point) {
 			Ok(blocks) => blocks,
 			Err(err) => {
 				debug!(
@@ -707,10 +709,6 @@ where
 		let merged = tokio_stream::StreamExt::merge(merged, stream_responses);
 		let stream = stream::once(futures::future::ready(initial)).chain(merged);
 
-		// 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
 	}
 }

From e4ccc032a1f5fe00fa8b2e765aed4e811ef5628a Mon Sep 17 00:00:00 2001
From: Alexandru Vasile <alexandru.vasile@parity.io>
Date: Tue, 16 Apr 2024 17:53:39 +0300
Subject: [PATCH 16/16] chainHead/tests: Util code to reduce testing scenarios

Signed-off-by: Alexandru Vasile <alexandru.vasile@parity.io>
---
 .../rpc-spec-v2/src/chain_head/tests.rs       | 346 +++++-------------
 1 file changed, 99 insertions(+), 247 deletions(-)

diff --git a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
index 3347d6d8ac6d..14f664858a0d 100644
--- a/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
+++ b/substrate/client/rpc-spec-v2/src/chain_head/tests.rs
@@ -187,6 +187,62 @@ async fn setup_api() -> (
 	(client, api, sub, sub_id, block)
 }
 
+async fn import_block(
+	mut client: Arc<Client<Backend>>,
+	parent_hash: <Block as BlockT>::Hash,
+	parent_number: u64,
+) -> Block {
+	let block = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(parent_hash)
+		.with_parent_block_number(parent_number)
+		.build()
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	client.import(BlockOrigin::Own, block.clone()).await.unwrap();
+	block
+}
+
+async fn import_best_block_with_tx(
+	mut client: Arc<Client<Backend>>,
+	parent_hash: <Block as BlockT>::Hash,
+	parent_number: u64,
+	tx: Transfer,
+) -> Block {
+	let mut block_builder = BlockBuilderBuilder::new(&*client)
+		.on_parent_block(parent_hash)
+		.with_parent_block_number(parent_number)
+		.build()
+		.unwrap();
+	block_builder.push_transfer(tx).unwrap();
+	let block = block_builder.build().unwrap().block;
+	client.import_as_best(BlockOrigin::Own, block.clone()).await.unwrap();
+	block
+}
+
+/// Check the subscription produces a new block and a best block event.
+///
+/// The macro is used instead of a fn to preserve the lines of code in case of panics.
+macro_rules! check_new_and_best_block_events {
+	($sub:expr, $block_hash:expr, $parent_hash:expr) => {
+		let event: FollowEvent<String> = get_next_event($sub).await;
+		let expected = FollowEvent::NewBlock(NewBlock {
+			block_hash: format!("{:?}", $block_hash),
+			parent_block_hash: format!("{:?}", $parent_hash),
+			new_runtime: None,
+			with_runtime: false,
+		});
+		assert_eq!(event, expected);
+
+		let event: FollowEvent<String> = get_next_event($sub).await;
+		let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
+			best_block_hash: format!("{:?}", $block_hash),
+		});
+		assert_eq!(event, expected);
+	};
+}
+
 #[tokio::test]
 async fn follow_subscription_produces_blocks() {
 	let builder = TestClientBuilder::new();
@@ -3649,7 +3705,7 @@ async fn chain_head_limit_reached() {
 async fn follow_unique_pruned_blocks() {
 	let builder = TestClientBuilder::new();
 	let backend = builder.backend();
-	let mut client = Arc::new(builder.build());
+	let client = Arc::new(builder.build());
 
 	let api = ChainHead::new(
 		client.clone(),
@@ -3704,206 +3760,46 @@ async fn follow_unique_pruned_blocks() {
 	// only blocks 5 and 4 are reported as pruned. This is because the block 2 was previously
 	// reported.
 
-	let block_1 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(client.chain_info().genesis_hash)
-		.with_parent_block_number(0)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_1_hash = block_1.hash();
-	client.import(BlockOrigin::Own, block_1.clone()).await.unwrap();
-
-	let block_2_f = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_1_hash)
-		.with_parent_block_number(1)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_2_f_hash = block_2_f.hash();
-	client.import(BlockOrigin::Own, block_2_f.clone()).await.unwrap();
-
-	let block_6 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_2_f_hash)
-		.with_parent_block_number(2)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_6_hash = block_6.hash();
-
-	client.import(BlockOrigin::Own, block_6.clone()).await.unwrap();
-
+	// Initial setup steps:
+	let block_1_hash =
+		import_block(client.clone(), client.chain_info().genesis_hash, 0).await.hash();
+	let block_2_f_hash = import_block(client.clone(), block_1_hash, 1).await.hash();
+	let block_6_hash = import_block(client.clone(), block_2_f_hash, 2).await.hash();
 	// Import block 2 as best on the fork.
-	let mut block_builder = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_1_hash)
-		.with_parent_block_number(1)
-		.build()
-		.unwrap();
-	// This push is required as otherwise block 3 has the same hash as block 2 and won't get
-	// imported
-	block_builder
-		.push_transfer(Transfer {
-			from: AccountKeyring::Alice.into(),
-			to: AccountKeyring::Ferdie.into(),
-			amount: 41,
-			nonce: 0,
-		})
-		.unwrap();
-	let block_2 = block_builder.build().unwrap().block;
-	let block_2_hash = block_2.header.hash();
-	client.import_as_best(BlockOrigin::Own, block_2.clone()).await.unwrap();
-
-	let block_3 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_2_hash)
-		.with_parent_block_number(2)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_3_hash = block_3.hash();
-	client.import(BlockOrigin::Own, block_3.clone()).await.unwrap();
+	let mut tx_alice_ferdie = Transfer {
+		from: AccountKeyring::Alice.into(),
+		to: AccountKeyring::Ferdie.into(),
+		amount: 41,
+		nonce: 0,
+	};
+	let block_2_hash =
+		import_best_block_with_tx(client.clone(), block_1_hash, 1, tx_alice_ferdie.clone())
+			.await
+			.hash();
 
+	let block_3_hash = import_block(client.clone(), block_2_hash, 2).await.hash();
 	// Fork block 4.
-	let mut block_builder = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_2_hash)
-		.with_parent_block_number(2)
-		.build()
-		.unwrap();
-	// This push is required as otherwise block 3 has the same hash as block 2 and won't get
-	// imported
-	block_builder
-		.push_transfer(Transfer {
-			from: AccountKeyring::Alice.into(),
-			to: AccountKeyring::Ferdie.into(),
-			amount: 41,
-			nonce: 1,
-		})
-		.unwrap();
-	let block_4 = block_builder.build().unwrap().block;
-	let block_4_hash = block_4.header.hash();
-	client.import_as_best(BlockOrigin::Own, block_4.clone()).await.unwrap();
-
-	let block_5 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_4_hash)
-		.with_parent_block_number(3)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_5_hash = block_5.hash();
-	client.import(BlockOrigin::Own, block_5.clone()).await.unwrap();
-
-	// Check block 1.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_1_hash),
-		parent_block_hash: format!("{:?}", finalized_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_1_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 2f.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_2_f_hash),
-		parent_block_hash: format!("{:?}", block_1_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_2_f_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 6.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_6_hash),
-		parent_block_hash: format!("{:?}", block_2_f_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_6_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 2, that we imported as custom best.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_2_hash),
-		parent_block_hash: format!("{:?}", block_1_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_2_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 3.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_3_hash),
-		parent_block_hash: format!("{:?}", block_2_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_3_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 4.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_4_hash),
-		parent_block_hash: format!("{:?}", block_2_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_4_hash),
-	});
-	assert_eq!(event, expected);
-
-	// Check block 5.
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_5_hash),
-		parent_block_hash: format!("{:?}", block_4_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_5_hash),
-	});
-	assert_eq!(event, expected);
+	tx_alice_ferdie.nonce = 1;
+	let block_4_hash = import_best_block_with_tx(client.clone(), block_2_hash, 2, tx_alice_ferdie)
+		.await
+		.hash();
+	let block_5_hash = import_block(client.clone(), block_4_hash, 3).await.hash();
+
+	// Check expected events generated by the setup.
+	{
+		// Check block 1 -> block 2f -> block 6.
+		check_new_and_best_block_events!(&mut sub, block_1_hash, finalized_hash);
+		check_new_and_best_block_events!(&mut sub, block_2_f_hash, block_1_hash);
+		check_new_and_best_block_events!(&mut sub, block_6_hash, block_2_f_hash);
+
+		// Check (block 1 ->) block 2 -> block 3.
+		check_new_and_best_block_events!(&mut sub, block_2_hash, block_1_hash);
+		check_new_and_best_block_events!(&mut sub, block_3_hash, block_2_hash);
+
+		// Check (block 1 -> block 2 ->) block 4 -> block 5.
+		check_new_and_best_block_events!(&mut sub, block_4_hash, block_2_hash);
+		check_new_and_best_block_events!(&mut sub, block_5_hash, block_4_hash);
+	}
 
 	// Finalize the block 6 from the fork.
 	client.finalize_block(block_6_hash, None).unwrap();
@@ -3933,31 +3829,9 @@ async fn follow_unique_pruned_blocks() {
 	let hash = format!("{:?}", block_2_hash);
 	let _res: () = api.call("chainHead_unstable_unpin", rpc_params![&sub_id, &hash]).await.unwrap();
 
-	// Check block 7.
-	let block_7 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_6_hash)
-		.with_parent_block_number(3)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_7_hash = block_7.hash();
-	client.import(BlockOrigin::Own, block_7.clone()).await.unwrap();
-
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_7_hash),
-		parent_block_hash: format!("{:?}", block_6_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_7_hash),
-	});
-	assert_eq!(event, expected);
+	// Import block 7 and check it.
+	let block_7_hash = import_block(client.clone(), block_6_hash, 3).await.hash();
+	check_new_and_best_block_events!(&mut sub, block_7_hash, block_6_hash);
 
 	// Finalize the block 7.
 	client.finalize_block(block_7_hash, None).unwrap();
@@ -3970,30 +3844,8 @@ async fn follow_unique_pruned_blocks() {
 	assert_eq!(event, expected);
 
 	// Check block 8.
-	let block_8 = BlockBuilderBuilder::new(&*client)
-		.on_parent_block(block_7_hash)
-		.with_parent_block_number(4)
-		.build()
-		.unwrap()
-		.build()
-		.unwrap()
-		.block;
-	let block_8_hash = block_8.hash();
-	client.import(BlockOrigin::Own, block_8.clone()).await.unwrap();
-
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::NewBlock(NewBlock {
-		block_hash: format!("{:?}", block_8_hash),
-		parent_block_hash: format!("{:?}", block_7_hash),
-		new_runtime: None,
-		with_runtime: false,
-	});
-	assert_eq!(event, expected);
-	let event: FollowEvent<String> = get_next_event(&mut sub).await;
-	let expected = FollowEvent::BestBlockChanged(BestBlockChanged {
-		best_block_hash: format!("{:?}", block_8_hash),
-	});
-	assert_eq!(event, expected);
+	let block_8_hash = import_block(client.clone(), block_7_hash, 4).await.hash();
+	check_new_and_best_block_events!(&mut sub, block_8_hash, block_7_hash);
 
 	// Finalize the block 8.
 	client.finalize_block(block_8_hash, None).unwrap();