Finality notification: Optimize calculation of stale heads

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]>,

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!()
 	}

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

client/consensus/babe/src/tests.rs

@@ -1018,12 +1018,12 @@ fn obsolete_blocks_aux_data_cleanup() {
 
 	// Create the following test scenario:
 	//
-	//               /-----B3 --- B4           ( < fork2 )
+	//  /--- --B3 --- B4                       ( < fork2 )
 	// G --- A1 --- A2 --- A3 --- A4           ( < fork1 )
 	//                      \-----C4 --- C5    ( < fork3 )
 
 	let fork1_hashes = propose_and_import_blocks_wrap(BlockId::Number(0), 4);
-	let fork2_hashes = propose_and_import_blocks_wrap(BlockId::Number(2), 2);
+	let fork2_hashes = propose_and_import_blocks_wrap(BlockId::Number(0), 2);
 	let fork3_hashes = propose_and_import_blocks_wrap(BlockId::Number(3), 2);
 
 	// Check that aux data is present for all but the genesis block.

client/db/src/lib.rs

@@ -627,6 +627,19 @@ impl<Block: BlockT> sc_client_api::blockchain::Backend<Block> for BlockchainDb<B
 		Ok(self.leaves.read().hashes())
 	}
 
+	fn displaced_leaves_after_finalizing(
+		&self,
+		block_number: NumberFor<Block>,
+	) -> ClientResult<Vec<Block::Hash>> {
+		Ok(self
+			.leaves
+			.read()
+			.displaced_by_finalize_height(block_number)
+			.leaves()
+			.cloned()
+			.collect::<Vec<_>>())
+	}
+
 	fn children(&self, parent_hash: Block::Hash) -> ClientResult<Vec<Block::Hash>> {
 		children::read_children(&*self.db, columns::META, meta_keys::CHILDREN_PREFIX, parent_hash)
 	}

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 });
 		}
 

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());
 }
 

primitives/blockchain/src/backend.rs

@@ -100,6 +100,14 @@ pub trait Backend<Block: BlockT>:
 	/// Results must be ordered best (longest, highest) chain first.
 	fn leaves(&self) -> Result<Vec<Block::Hash>>;
 
+	/// Returns displaced leaves after the given block would be finalized.
+	///
+	/// The returned leaves do not contain the leaves from the same height as `block_number`.
+	fn displaced_leaves_after_finalizing(
+		&self,
+		block_number: NumberFor<Block>,
+	) -> Result<Vec<Block::Hash>>;
+
 	/// Return hashes of all blocks that are children of the block with `parent_hash`.
 	fn children(&self, parent_hash: Block::Hash) -> Result<Vec<Block::Hash>>;
 

primitives/blockchain/src/header_metadata.rs

@@ -201,6 +201,11 @@ impl<Block: BlockT> TreeRoute<Block> {
 	pub fn enacted(&self) -> &[HashAndNumber<Block>] {
 		&self.route[self.pivot + 1..]
 	}
+
+	/// Returns the last block.
+	pub fn last(&self) -> Option<&HashAndNumber<Block>> {
+		self.route.last()
+	}
 }
 
 /// Handles header metadata: hash, number, parent hash, etc.