Finality notification: Optimize calculation of stale heads (#11200)

* Finality notification: Optimize calculation of stale heads

While looking into some problem on Versi where a collator seemed to be stuck. I found out that it
was not stuck but there was a huge gap between last finalized and best block. This lead to a lot
leaves and it was basically trapped inside some loop of reading block headers from the db to find
the stale heads. While looking into this I found out that `leaves` already supports the feature to
give us the stale heads relative easily. However, the semantics change a little bit. Instead of
returning all stale heads of blocks that are not reachable anymore after finalizing a block, we
currently only return heads with a number lower than the finalized block. This should be no problem,
because these other leaves that are stale will be returned later when a block gets finalized which
number is bigger than the block number of these leaves.

While doing that, I also changed `tree_route` of the `FinalityNotification` to include the
`old_finalized`. Based on the comment I assumed that this was already part of it. However, if
wanted, I can revert this change.

* FMT

* Update client/service/src/client/client.rs

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

* Do not include the last finalized block

* Rename function

* FMT

* Fix tests

* Update figure

Co-authored-by: André Silva <123550+andresilva@users.noreply.github.com>

substrate/client/service/src/client/client.rs

@@ -899,36 +899,25 @@ where
 			let finalized =
 				route_from_finalized.enacted().iter().map(|elem| elem.hash).collect::<Vec<_>>();
 
-			let last_finalized_number = self
-				.backend
-				.blockchain()
-				.number(last_finalized)?
-				.expect("Previous finalized block expected to be onchain; qed");
-			let mut stale_heads = Vec::new();
-			for head in self.backend.blockchain().leaves()? {
-				let route_from_finalized =
-					sp_blockchain::tree_route(self.backend.blockchain(), block, head)?;
-				let retracted = route_from_finalized.retracted();
-				let pivot = route_from_finalized.common_block();
-				// It is not guaranteed that `backend.blockchain().leaves()` doesn't return
-				// heads that were in a stale state before this finalization and thus already
-				// included in previous notifications. We want to skip such heads.
-				// Given the "route" from the currently finalized block to the head under
-				// analysis, the condition for it to be added to the new stale heads list is:
-				// `!retracted.is_empty() && last_finalized_number <= pivot.number`
-				// 1. "route" has some "retractions".
-				// 2. previously finalized block number is not greater than the "route" pivot:
-				//    - if `last_finalized_number <= pivot.number` then this is a new stale head;
-				//    - else the stale head was already included by some previous finalization.
-				if !retracted.is_empty() && last_finalized_number <= pivot.number {
-					stale_heads.push(head);
-				}
-			}
+			let block_number = route_from_finalized
+				.last()
+				.expect(
+					"The block to finalize is always the latest \
+						block in the route to the finalized block; qed",
+				)
+				.number;
+
+			// The stale heads are the leaves that will be displaced after the
+			// block is finalized.
+			let stale_heads =
+				self.backend.blockchain().displaced_leaves_after_finalizing(block_number)?;
+
 			let header = self
 				.backend
 				.blockchain()
 				.header(BlockId::Hash(block))?
-				.expect("Finalized block expected to be onchain; qed");
+				.expect("Block to finalize expected to be onchain; qed");
+
 			operation.notify_finalized = Some(FinalizeSummary { header, finalized, stale_heads });
 		}
 

substrate/client/service/test/src/client/mod.rs

@@ -1011,12 +1011,16 @@ fn finalizing_diverged_block_should_trigger_reorg() {
 
 	assert_eq!(client.chain_info().best_hash, b3.hash());
 
-	finality_notification_check(&mut finality_notifications, &[b1.hash()], &[a2.hash()]);
+	ClientExt::finalize_block(&client, BlockId::Hash(b3.hash()), None).unwrap();
+
+	finality_notification_check(&mut finality_notifications, &[b1.hash()], &[]);
+	finality_notification_check(&mut finality_notifications, &[b2.hash(), b3.hash()], &[a2.hash()]);
 	assert!(finality_notifications.try_next().is_err());
 }
 
 #[test]
 fn finality_notifications_content() {
+	sp_tracing::try_init_simple();
 	let (mut client, _select_chain) = TestClientBuilder::new().build_with_longest_chain();
 
 	//               -> D3 -> D4
@@ -1110,12 +1114,8 @@ fn finality_notifications_content() {
 	// Import and finalize D4
 	block_on(client.import_as_final(BlockOrigin::Own, d4.clone())).unwrap();
 
-	finality_notification_check(
-		&mut finality_notifications,
-		&[a1.hash(), a2.hash()],
-		&[c1.hash(), b2.hash()],
-	);
-	finality_notification_check(&mut finality_notifications, &[d3.hash(), d4.hash()], &[a3.hash()]);
+	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[c1.hash()]);
+	finality_notification_check(&mut finality_notifications, &[d3.hash(), d4.hash()], &[b2.hash()]);
 	assert!(finality_notifications.try_next().is_err());
 }
 
@@ -1214,7 +1214,7 @@ fn doesnt_import_blocks_that_revert_finality() {
 
 	//    -> C1
 	//   /
-	// G -> A1 -> A2
+	// G -> A1 -> A2 -> A3
 	//   \
 	//    -> B1 -> B2 -> B3
 
@@ -1294,7 +1294,19 @@ fn doesnt_import_blocks_that_revert_finality() {
 
 	assert_eq!(import_err.to_string(), expected_err.to_string());
 
-	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[b2.hash()]);
+	let a3 = client
+		.new_block_at(&BlockId::Hash(a2.hash()), Default::default(), false)
+		.unwrap()
+		.build()
+		.unwrap()
+		.block;
+	block_on(client.import(BlockOrigin::Own, a3.clone())).unwrap();
+	ClientExt::finalize_block(&client, BlockId::Hash(a3.hash()), None).unwrap();
+
+	finality_notification_check(&mut finality_notifications, &[a1.hash(), a2.hash()], &[]);
+
+	finality_notification_check(&mut finality_notifications, &[a3.hash()], &[b2.hash()]);
+
 	assert!(finality_notifications.try_next().is_err());
 }