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

@@ -1369,6 +1369,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>,
@@ -1377,6 +1379,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();
@@ -1389,6 +1392,7 @@ impl<Block: BlockT> Backend<Block> {
 			hash,
 			with_state,
 			current_transaction_justifications,
+			remove_displaced,
 		)?;
 
 		if let Some(justification) = justification {
@@ -1466,7 +1470,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,
@@ -1475,6 +1480,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();
@@ -1654,6 +1660,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.
@@ -1779,16 +1786,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();
 
@@ -1813,13 +1822,19 @@ 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(),
+			));
+
+			if !matches!(self.blocks_pruning, BlocksPruning::KeepAll) {
+				self.prune_displaced_branches(transaction, &new_displaced)?;
+			}
+		}
 
-		self.prune_blocks(transaction, f_num, &new_displaced, current_transaction_justifications)?;
+		self.prune_blocks(transaction, f_num, current_transaction_justifications)?;
 
 		Ok(())
 	}
@@ -1828,39 +1843,29 @@ impl<Block: BlockT> Backend<Block> {
 		&self,
 		transaction: &mut Transaction<DbHash>,
 		finalized_number: NumberFor<Block>,
-		displaced: &DisplacedLeavesAfterFinalization<Block>,
 		current_transaction_justifications: &mut HashMap<Block::Hash, Justification>,
 	) -> ClientResult<()> {
-		match self.blocks_pruning {
-			BlocksPruning::KeepAll => {},
-			BlocksPruning::Some(blocks_pruning) => {
-				// Always keep the last finalized block
-				let keep = std::cmp::max(blocks_pruning, 1);
-				if finalized_number >= keep.into() {
-					let number = finalized_number.saturating_sub(keep.into());
-
-					// Before we prune a block, check if it is pinned
-					if let Some(hash) = self.blockchain.hash(number)? {
-						self.blockchain.insert_persisted_body_if_pinned(hash)?;
-
-						// If the block was finalized in this transaction, it will not be in the db
-						// yet.
-						if let Some(justification) =
-							current_transaction_justifications.remove(&hash)
-						{
-							self.blockchain.insert_justifications_if_pinned(hash, justification);
-						} else {
-							self.blockchain.insert_persisted_justifications_if_pinned(hash)?;
-						}
-					};
+		if let BlocksPruning::Some(blocks_pruning) = self.blocks_pruning {
+			// Always keep the last finalized block
+			let keep = std::cmp::max(blocks_pruning, 1);
+			if finalized_number >= keep.into() {
+				let number = finalized_number.saturating_sub(keep.into());
+
+				// Before we prune a block, check if it is pinned
+				if let Some(hash) = self.blockchain.hash(number)? {
+					self.blockchain.insert_persisted_body_if_pinned(hash)?;
+
+					// If the block was finalized in this transaction, it will not be in the db
+					// yet.
+					if let Some(justification) = current_transaction_justifications.remove(&hash) {
+						self.blockchain.insert_justifications_if_pinned(hash, justification);
+					} else {
+						self.blockchain.insert_persisted_justifications_if_pinned(hash)?;
+					}
+				};
 
-					self.prune_block(transaction, BlockId::<Block>::number(number))?;
-				}
-				self.prune_displaced_branches(transaction, displaced)?;
-			},
-			BlocksPruning::KeepFinalized => {
-				self.prune_displaced_branches(transaction, displaced)?;
-			},
+				self.prune_block(transaction, BlockId::<Block>::number(number))?;
+			}
 		}
 		Ok(())
 	}
@@ -2121,6 +2126,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)?;