Optimize pruning during batch finalization of blocks

substrate/client/api/src/backend.rs

@@ -217,7 +217,8 @@ pub trait BlockImportOperation<Block: BlockT> {
 	where
 		I: IntoIterator<Item = (Vec<u8>, Option<Vec<u8>>)>;
 
-	/// Mark a block as finalized.
+	/// Mark a block as finalized, if multiple blocks are finalized in the same operation then they
+	/// must be marked in ascending order.
 	fn mark_finalized(
 		&mut self,
 		hash: Block::Hash,

substrate/client/db/src/lib.rs

@@ -1357,6 +1357,8 @@ impl<Block: BlockT> Backend<Block> {
 		Ok(())
 	}
 
+	/// `remove_displaced` can be set to `false` if this is not the last of many subsequent calls
+	/// for performance reasons.
 	fn finalize_block_with_transaction(
 		&self,
 		transaction: &mut Transaction<DbHash>,
@@ -1365,6 +1367,7 @@ impl<Block: BlockT> Backend<Block> {
 		last_finalized: Option<Block::Hash>,
 		justification: Option<Justification>,
 		current_transaction_justifications: &mut HashMap<Block::Hash, Justification>,
+		remove_displaced: bool,
 	) -> ClientResult<MetaUpdate<Block>> {
 		// TODO: ensure best chain contains this block.
 		let number = *header.number();
@@ -1377,6 +1380,7 @@ impl<Block: BlockT> Backend<Block> {
 			hash,
 			with_state,
 			current_transaction_justifications,
+			remove_displaced,
 		)?;
 
 		if let Some(justification) = justification {
@@ -1454,7 +1458,8 @@ impl<Block: BlockT> Backend<Block> {
 
 		let mut current_transaction_justifications: HashMap<Block::Hash, Justification> =
 			HashMap::new();
-		for (block_hash, justification) in operation.finalized_blocks {
+		let mut finalized_blocks = operation.finalized_blocks.into_iter().peekable();
+		while let Some((block_hash, justification)) = finalized_blocks.next() {
 			let block_header = self.blockchain.expect_header(block_hash)?;
 			meta_updates.push(self.finalize_block_with_transaction(
 				&mut transaction,
@@ -1463,6 +1468,7 @@ impl<Block: BlockT> Backend<Block> {
 				Some(last_finalized_hash),
 				justification,
 				&mut current_transaction_justifications,
+				finalized_blocks.peek().is_none(),
 			)?);
 			last_finalized_hash = block_hash;
 			last_finalized_num = *block_header.number();
@@ -1642,6 +1648,7 @@ impl<Block: BlockT> Backend<Block> {
 					hash,
 					operation.commit_state,
 					&mut current_transaction_justifications,
+					true,
 				)?;
 			} else {
 				// canonicalize blocks which are old enough, regardless of finality.
@@ -1766,16 +1773,18 @@ impl<Block: BlockT> Backend<Block> {
 		Ok(())
 	}
 
-	// write stuff to a transaction after a new block is finalized.
-	// this canonicalizes finalized blocks. Fails if called with a block which
-	// was not a child of the last finalized block.
+	// Write stuff to a transaction after a new block is finalized. This canonicalizes finalized
+	// blocks. Fails if called with a block which was not a child of the last finalized block.
+	/// `remove_displaced` can be set to `false` if this is not the last of many subsequent calls
+	/// for performance reasons.
 	fn note_finalized(
 		&self,
 		transaction: &mut Transaction<DbHash>,
 		f_header: &Block::Header,
 		f_hash: Block::Hash,
 		with_state: bool,
 		current_transaction_justifications: &mut HashMap<Block::Hash, Justification>,
+		remove_displaced: bool,
 	) -> ClientResult<()> {
 		let f_num = *f_header.number();
 
@@ -1800,13 +1809,20 @@ impl<Block: BlockT> Backend<Block> {
 			apply_state_commit(transaction, commit);
 		}
 
-		let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
+		if remove_displaced {
+			let new_displaced = self.blockchain.displaced_leaves_after_finalizing(f_hash, f_num)?;
 
-		self.blockchain.leaves.write().remove_displaced_leaves(FinalizationOutcome::new(
-			new_displaced.displaced_leaves.iter().copied(),
-		));
+			self.blockchain.leaves.write().remove_displaced_leaves(FinalizationOutcome::new(
+				new_displaced.displaced_leaves.iter().copied(),
+			));
 
-		self.prune_blocks(transaction, f_num, &new_displaced, current_transaction_justifications)?;
+			self.prune_blocks(
+				transaction,
+				f_num,
+				&new_displaced,
+				current_transaction_justifications,
+			)?;
+		}
 
 		Ok(())
 	}
@@ -2108,6 +2124,7 @@ impl<Block: BlockT> sc_client_api::backend::Backend<Block> for Backend<Block> {
 			None,
 			justification,
 			&mut current_transaction_justifications,
+			true,
 		)?;
 
 		self.storage.db.commit(transaction)?;