Change forks pruning algorithm.

substrate/client/api/src/in_mem.rs

@@ -419,20 +419,6 @@ 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

@@ -49,7 +49,7 @@ pub struct FinalizationOutcome<H, N> {
 	removed: BTreeMap<Reverse<N>, Vec<H>>,
 }
 
-impl<H, N: Ord> FinalizationOutcome<H, N> {
+impl<H: Copy, N: Ord> FinalizationOutcome<H, N> {
 	/// Merge with another. This should only be used for displaced items that
 	/// are produced within one transaction of each other.
 	pub fn merge(&mut self, mut other: Self) {
@@ -63,6 +63,21 @@ impl<H, N: Ord> FinalizationOutcome<H, N> {
 	pub fn leaves(&self) -> impl Iterator<Item = &H> {
 		self.removed.values().flatten()
 	}
+
+	/// Constructor
+	pub fn new(new_displaced: impl Iterator<Item = (H, N)>) -> Self {
+		let mut removed = BTreeMap::<Reverse<N>, Vec<H>>::new();
+		for (hash, number) in new_displaced {
+			removed
+				.entry(Reverse(number))
+				.and_modify(|hashes| {
+					hashes.push(hash);
+				})
+				.or_insert(vec![hash]);
+		}
+
+		FinalizationOutcome { removed }
+	}
 }
 
 /// list of leaf hashes ordered by number (descending).
@@ -151,39 +166,12 @@ where
 		Some(RemoveOutcome { inserted, removed: LeafSetItem { hash, number } })
 	}
 
-	/// Note a block height finalized, displacing all leaves with number less than the finalized
-	/// block's.
-	///
-	/// Although it would be more technically correct to also prune out leaves at the
-	/// same number as the finalized block, but with different hashes, the current behavior
-	/// is simpler and our assumptions about how finalization works means that those leaves
-	/// will be pruned soon afterwards anyway.
-	pub fn finalize_height(&mut self, number: N) -> FinalizationOutcome<H, N> {
-		let boundary = if number == N::zero() {
-			return FinalizationOutcome { removed: BTreeMap::new() }
-		} else {
-			number - N::one()
-		};
-
-		let below_boundary = self.storage.split_off(&Reverse(boundary));
-		FinalizationOutcome { removed: 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) -> FinalizationOutcome<H, N> {
-		let boundary = if number == N::zero() {
-			return FinalizationOutcome { removed: BTreeMap::new() }
-		} else {
-			number - N::one()
-		};
-
-		let below_boundary = self.storage.range(&Reverse(boundary)..);
-		FinalizationOutcome {
-			removed: below_boundary.map(|(k, v)| (k.clone(), v.clone())).collect(),
+	/// Remove all leaves displaced by the last block finalization.
+	pub fn remove_displaced_leaves(&mut self, displaced_leaves: &FinalizationOutcome<H, N>) {
+		for (number, hashes) in &displaced_leaves.removed {
+			for hash in hashes.iter() {
+				self.remove_leaf(number, hash);
+			}
 		}
 	}
 

substrate/client/db/src/lib.rs

@@ -747,19 +747,6 @@ 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)
 	}
@@ -1813,12 +1800,16 @@ impl<Block: BlockT> Backend<Block> {
 			apply_state_commit(transaction, commit);
 		}
 
-		let new_displaced = self.blockchain.leaves.write().finalize_height(f_num);
+		let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
+		let finalization_outcome = FinalizationOutcome::new(new_displaced.into_iter());
+
+		self.blockchain.leaves.write().remove_displaced_leaves(&finalization_outcome);
+
 		self.prune_blocks(
 			transaction,
 			f_num,
 			f_hash,
-			&new_displaced,
+			&finalization_outcome,
 			current_transaction_justifications,
 		)?;
 
@@ -3190,6 +3181,9 @@ pub(crate) mod tests {
 
 	#[test]
 	fn test_leaves_pruned_on_finality() {
+		//   / 1b - 2b - 3b
+		// 0 - 1a - 2a
+		//   \ 1c
 		let backend: Backend<Block> = Backend::new_test(10, 10);
 		let block0 = insert_header(&backend, 0, Default::default(), None, Default::default());
 
@@ -3201,18 +3195,16 @@ pub(crate) mod tests {
 
 		let block2_a = insert_header(&backend, 2, block1_a, None, Default::default());
 		let block2_b = insert_header(&backend, 2, block1_b, None, Default::default());
-		let block2_c = insert_header(&backend, 2, block1_b, None, [1; 32].into());
 
-		assert_eq!(
-			backend.blockchain().leaves().unwrap(),
-			vec![block2_a, block2_b, block2_c, block1_c]
-		);
+		let block3_b = insert_header(&backend, 3, block2_b, None, [3; 32].into());
+
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block3_b, block2_a, block1_c]);
 
 		backend.finalize_block(block1_a, None).unwrap();
 		backend.finalize_block(block2_a, None).unwrap();
 
-		// leaves at same height stay. Leaves at lower heights pruned.
-		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block2_a, block2_b, block2_c]);
+		// All leaves are pruned that are known to not belong to canonical branch
+		assert_eq!(backend.blockchain().leaves().unwrap(), vec![block2_a]);
 	}
 
 	#[test]

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

@@ -978,8 +978,13 @@ where
 
 			// 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 stale_heads = self
+				.backend
+				.blockchain()
+				.displaced_leaves_after_finalizing(hash, block_number)?
+				.into_iter()
+				.map(|(h, _)| h)
+				.collect();
 
 			let header = self
 				.backend

substrate/primitives/blockchain/src/backend.rs

@@ -21,7 +21,7 @@ use log::warn;
 use parking_lot::RwLock;
 use sp_runtime::{
 	generic::BlockId,
-	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating},
+	traits::{Block as BlockT, Header as HeaderT, NumberFor, Saturating, Zero},
 	Justifications,
 };
 use std::collections::btree_set::BTreeSet;
@@ -172,14 +172,6 @@ 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>>;
 
@@ -255,6 +247,58 @@ pub trait Backend<Block: BlockT>:
 	}
 
 	fn block_indexed_body(&self, hash: Block::Hash) -> Result<Option<Vec<Vec<u8>>>>;
+
+	/// Returns all leaves that will be displaced after the block finalization.
+	fn displaced_leaves_after_finalizing(
+		&self,
+		finalized_block_hash: Block::Hash,
+		finalized_block_number: NumberFor<Block>,
+	) -> std::result::Result<Vec<(Block::Hash, NumberFor<Block>)>, Error> {
+		if finalized_block_number == Zero::zero() {
+			return Ok(Vec::new())
+		}
+
+		let mut displaced_leaves = Vec::new();
+
+		let finalized_block_header = self.expect_header(finalized_block_hash)?;
+		// For each leaf determine whether it belongs to a non-canonical branch.
+		for leaf_hash in self.leaves()? {
+			let mut fork_block_header = self.expect_header(leaf_hash)?;
+
+			let leaf_number = *fork_block_header.number();
+
+			let mut needs_pruning = false;
+			// All leaves will eventually have as ancestor either the finalized block or its
+			// parent. All other forks were cleared on the previous block finalization.
+			loop {
+				// The block ends up in the finalized block. All forks are valid at this point.
+				if fork_block_header.hash() == finalized_block_hash {
+					break
+				}
+
+				// The block ends up in the parent block of the finalized block. It's a stale fork.
+				if fork_block_header.hash() == *finalized_block_header.parent_hash() {
+					needs_pruning = true;
+					break
+				}
+
+				if let Some(parent_header) = self.header(*fork_block_header.parent_hash())? {
+					fork_block_header = parent_header;
+				} else {
+					// Sometimes routes can't be calculated. E.g. after warp sync.
+					needs_pruning = true;
+					break
+				}
+			}
+
+			// Fork ended up in the parent of the finalized block.
+			if needs_pruning {
+				displaced_leaves.push((leaf_hash, leaf_number));
+			}
+		}
+
+		Ok(displaced_leaves)
+	}
 }
 
 /// Blockchain info