finalization: Skip tree route calculation if no forks present

substrate/client/db/src/lib.rs

@@ -2547,7 +2547,7 @@ pub(crate) mod tests {
 		backend::{Backend as BTrait, BlockImportOperation as Op},
 		blockchain::Backend as BLBTrait,
 	};
-	use sp_blockchain::{lowest_common_ancestor, tree_route};
+	use sp_blockchain::{lowest_common_ancestor, lowest_common_ancestor_multiblock, tree_route};
 	use sp_core::H256;
 	use sp_runtime::{
 		testing::{Block as RawBlock, ExtrinsicWrapper, Header},
@@ -3108,6 +3108,111 @@ pub(crate) mod tests {
 		}
 	}
 
+	#[test]
+	fn lowest_common_ancestors_multiblock_works() {
+		let backend = Backend::<Block>::new_test(1000, 100);
+		let blockchain = backend.blockchain();
+		let block0 = insert_header(&backend, 0, Default::default(), None, Default::default());
+
+		// fork from genesis: 3 prong.
+		let a1 = insert_header(&backend, 1, block0, None, Default::default());
+		let a2 = insert_header(&backend, 2, a1, None, Default::default());
+		let a3 = insert_header(&backend, 3, a2, None, Default::default());
+
+		// fork from genesis: 2 prong.
+		let b1 = insert_header(&backend, 1, block0, None, H256::from([1; 32]));
+		let b2 = insert_header(&backend, 2, b1, None, Default::default());
+
+		// fork from b2.
+		let c1 = insert_header(&backend, 3, b2, None, H256::from([2; 32]));
+		let c2 = insert_header(&backend, 4, c1, None, Default::default());
+
+		// fork from b1.
+		let d1 = insert_header(&backend, 2, b1, None, H256::from([3; 32]));
+		let d2 = insert_header(&backend, 3, d1, None, Default::default());
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a3, b2]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, block0);
+			assert_eq!(lca.number, 0);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a1, a3]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a1);
+			assert_eq!(lca.number, 1);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a3, a1]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a1);
+			assert_eq!(lca.number, 1);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a2, a3]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a2);
+			assert_eq!(lca.number, 2);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a2, a1]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a1);
+			assert_eq!(lca.number, 1);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a2, a2]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a2);
+			assert_eq!(lca.number, 2);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a3, d2, c2])
+				.unwrap()
+				.unwrap();
+
+			assert_eq!(lca.hash, block0);
+			assert_eq!(lca.number, 0);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![c2, d2, b2])
+				.unwrap()
+				.unwrap();
+
+			assert_eq!(lca.hash, b1);
+			assert_eq!(lca.number, 1);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a1, a2, a3])
+				.unwrap()
+				.unwrap();
+
+			assert_eq!(lca.hash, a1);
+			assert_eq!(lca.number, 1);
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![]);
+
+			assert_eq!(true, matches!(lca, Ok(None)));
+		}
+
+		{
+			let lca = lowest_common_ancestor_multiblock(blockchain, vec![a1]).unwrap().unwrap();
+
+			assert_eq!(lca.hash, a1);
+			assert_eq!(lca.number, 1);
+		}
+	}
+
 	#[test]
 	fn test_tree_route_regression() {
 		// NOTE: this is a test for a regression introduced in #3665, the result

substrate/primitives/blockchain/src/backend.rs

@@ -21,16 +21,15 @@ use log::warn;
 use parking_lot::RwLock;
 use sp_runtime::{
 	generic::BlockId,
-	traits::{Block as BlockT, Header as HeaderT, NumberFor, Zero},
+	traits::{Block as BlockT, CheckedSub, Header as HeaderT, NumberFor, Zero},
 	Justifications,
 };
 use std::collections::{btree_map::BTreeMap, btree_set::BTreeSet};
 
-use crate::header_metadata::HeaderMetadata;
-
 use crate::{
 	error::{Error, Result},
-	tree_route, TreeRoute,
+	header_metadata::{self, HeaderMetadata},
+	lowest_common_ancestor_multiblock, tree_route, TreeRoute,
 };
 
 /// Blockchain database header backend. Does not perform any validation.
@@ -229,10 +228,37 @@ pub trait Backend<Block: BlockT>:
 	) -> std::result::Result<DisplacedLeavesAfterFinalization<Block>, Error> {
 		let mut result = DisplacedLeavesAfterFinalization::default();
 
-		if finalized_block_number == Zero::zero() {
+		let leaves = self.leaves()?;
+
+		// If we have only one leaf there are no forks, and we can return early.
+		if finalized_block_number == Zero::zero() || leaves.len() == 1 {
 			return Ok(result)
 		}
 
+		let first_leaf = leaves[0];
+		let leaf_block_header = self.expect_header(first_leaf)?;
+
+		// If the distance between the leafs and the finalized block is large,  calculating
+		// tree routes can be very expensive. In that case, we will try to find the
+		// lowest common ancestor between all the leaves. The assumption here is that the forks are
+		// close to the tip and not long. So the LCA can be computed from the header cache. If the
+		// LCA is above the finalized block, we know that there are no displaced leaves by the
+		// finalization.
+		if leaf_block_header
+			.number()
+			.checked_sub(&finalized_block_number)
+			.unwrap_or(0u32.into()) >
+			header_metadata::LRU_CACHE_SIZE.into()
+		{
+			if let Some(lca) = lowest_common_ancestor_multiblock(self, leaves)? {
+				if lca.number > finalized_block_number {
+					return Ok(result)
+				} else {
+					log::warn!("The distance between leafs and finalized block is large. Finalization can take a long time.");
+				}
+			};
+		}
+
 		// For each leaf determine whether it belongs to a non-canonical branch.
 		for leaf_hash in self.leaves()? {
 			let leaf_block_header = self.expect_header(leaf_hash)?;

substrate/primitives/blockchain/src/header_metadata.rs

@@ -23,7 +23,7 @@ use schnellru::{ByLength, LruMap};
 use sp_runtime::traits::{Block as BlockT, Header, NumberFor, One};
 
 /// Set to the expected max difference between `best` and `finalized` blocks at sync.
-const LRU_CACHE_SIZE: u32 = 5_000;
+pub(crate) const LRU_CACHE_SIZE: u32 = 5_000;
 
 /// Get lowest common ancestor between two blocks in the tree.
 ///
@@ -96,6 +96,27 @@ pub fn lowest_common_ancestor<Block: BlockT, T: HeaderMetadata<Block> + ?Sized>(
 	Ok(HashAndNumber { hash: header_one.hash, number: header_one.number })
 }
 
+/// Get lowest common ancestor between multiple blocks.
+pub fn lowest_common_ancestor_multiblock<Block: BlockT, T: HeaderMetadata<Block> + ?Sized>(
+	backend: &T,
+	hashes: Vec<Block::Hash>,
+) -> Result<Option<HashAndNumber<Block>>, T::Error> {
+	// Ensure the list of hashes is not empty
+	if hashes.is_empty() {
+		return Ok(None);
+	}
+
+	// Start with the first hash as the initial LCA
+	let cached = backend.header_metadata(hashes[0])?;
+	let mut lca = HashAndNumber { number: cached.number, hash: cached.hash };
+	for hash in hashes.iter().skip(1) {
+		// Calculate the LCA of the current LCA and the next hash
+		lca = lowest_common_ancestor(backend, lca.hash, *hash)?;
+	}
+
+	Ok(Some(lca))
+}
+
 /// Compute a tree-route between two blocks. See tree-route docs for more details.
 pub fn tree_route<Block: BlockT, T: HeaderMetadata<Block> + ?Sized>(
 	backend: &T,