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/api/src/client.rs

@@ -307,6 +307,8 @@ pub struct FinalityNotification<Block: BlockT> {
 	/// Finalized block header.
 	pub header: Block::Header,
 	/// Path from the old finalized to new finalized parent (implicitly finalized blocks).
+	///
+	/// This maps to the range `(old_finalized, new_finalized]`.
 	pub tree_route: Arc<[Block::Hash]>,
 	/// Stale branches heads.
 	pub stale_heads: Arc<[Block::Hash]>,

substrate/client/api/src/in_mem.rs

@@ -444,6 +444,20 @@ impl<Block: BlockT> blockchain::Backend<Block> for Blockchain<Block> {
 		Ok(self.storage.read().leaves.hashes())
 	}
 
+	fn displaced_leaves_after_finalizing(
+		&self,
+		block_number: NumberFor<Block>,
+	) -> sp_blockchain::Result<Vec<Block::Hash>> {
+		Ok(self
+			.storage
+			.read()
+			.leaves
+			.displaced_by_finalize_height(block_number)
+			.leaves()
+			.cloned()
+			.collect::<Vec<_>>())
+	}
+
 	fn children(&self, _parent_hash: Block::Hash) -> sp_blockchain::Result<Vec<Block::Hash>> {
 		unimplemented!()
 	}

substrate/client/api/src/leaves.rs

@@ -56,7 +56,7 @@ impl<H, N: Ord> FinalizationDisplaced<H, N> {
 	}
 
 	/// Iterate over all displaced leaves.
-	pub fn leaves(&self) -> impl IntoIterator<Item = &H> {
+	pub fn leaves(&self) -> impl Iterator<Item = &H> {
 		self.leaves.values().flatten()
 	}
 }
@@ -145,6 +145,24 @@ where
 		FinalizationDisplaced { leaves: below_boundary }
 	}
 
+	/// The same as [`Self::finalize_height`], but it only simulates the operation.
+	///
+	/// This means that no changes are done.
+	///
+	/// Returns the leaves that would be displaced by finalizing the given block.
+	pub fn displaced_by_finalize_height(&self, number: N) -> FinalizationDisplaced<H, N> {
+		let boundary = if number == N::zero() {
+			return FinalizationDisplaced { leaves: BTreeMap::new() }
+		} else {
+			number - N::one()
+		};
+
+		let below_boundary = self.storage.range(&Reverse(boundary)..);
+		FinalizationDisplaced {
+			leaves: below_boundary.map(|(k, v)| (k.clone(), v.clone())).collect(),
+		}
+	}
+
 	/// Undo all pending operations.
 	///
 	/// This returns an `Undo` struct, where any