BABE: Fix aux data cleaning (#11263)

With the latest optimizations of the `FinalityNotification` generation, the aux data pruning started
to print a warning. The problem here was that we printed a warning and stopped the adding of blocks
to prune when we hit the `heigh_limit`. This is now wrong, as we could for example have two 512 long
forks and then we start finalizing one of them. The second fork head would be part of the stale
heads at some point (in the current implementation when we finalize second fork head number + 1),
but then we would actually need to go back into the past than `heigh_limit` (which was actually
last_finalized - 1). We now go back until we reach the canonical chain.

Also fixed some wrong comment that was added by be about the content of the `finalized` blocks in
the `FinalityNotification`.

substrate/client/api/src/client.rs

@@ -308,7 +308,7 @@ pub struct FinalityNotification<Block: BlockT> {
 	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]`.
+	/// This maps to the range `(old_finalized, new_finalized)`.
 	pub tree_route: Arc<[Block::Hash]>,
 	/// Stale branches heads.
 	pub stale_heads: Arc<[Block::Hash]>,

substrate/client/consensus/babe/src/lib.rs

@@ -541,49 +541,66 @@ where
 // Remove obsolete block's weight data by leveraging finality notifications.
 // This includes data for all finalized blocks (excluding the most recent one)
 // and all stale branches.
-fn aux_storage_cleanup<C: HeaderMetadata<Block>, Block: BlockT>(
+fn aux_storage_cleanup<C: HeaderMetadata<Block> + HeaderBackend<Block>, Block: BlockT>(
 	client: &C,
 	notification: &FinalityNotification<Block>,
 ) -> AuxDataOperations {
 	let mut aux_keys = HashSet::new();
 
-	// Cleans data for finalized block's ancestors down to, and including, the previously
-	// finalized one.
-
-	let first_new_finalized = notification.tree_route.get(0).unwrap_or(&notification.hash);
-	match client.header_metadata(*first_new_finalized) {
+	let first = notification.tree_route.first().unwrap_or(&notification.hash);
+	match client.header_metadata(*first) {
 		Ok(meta) => {
 			aux_keys.insert(aux_schema::block_weight_key(meta.parent));
 		},
-		Err(err) => {
-			warn!(target: "babe", "header lookup fail while cleaning data for block {}: {}", first_new_finalized.to_string(), err.to_string());
-		},
+		Err(err) => warn!(
+			target: "babe",
+			"Failed to lookup metadata for block `{:?}`: {}",
+			first,
+			err,
+		),
 	}
 
-	aux_keys.extend(notification.tree_route.iter().map(aux_schema::block_weight_key));
+	// Cleans data for finalized block's ancestors
+	aux_keys.extend(
+		notification
+			.tree_route
+			.iter()
+			// Ensure we don't prune latest finalized block.
+			// This should not happen, but better be safe than sorry!
+			.filter(|h| **h != notification.hash)
+			.map(aux_schema::block_weight_key),
+	);
 
 	// Cleans data for stale branches.
 
-	// A safenet in case of malformed notification.
-	let height_limit = notification.header.number().saturating_sub(
-		notification.tree_route.len().saturated_into::<NumberFor<Block>>() + One::one(),
-	);
 	for head in notification.stale_heads.iter() {
 		let mut hash = *head;
-		// Insert stale blocks hashes until canonical chain is not reached.
-		// Soon or late we should hit an element already present within the `aux_keys` set.
+		// Insert stale blocks hashes until canonical chain is reached.
+		// If we reach a block that is already part of the `aux_keys` we can stop the processing the
+		// head.
 		while aux_keys.insert(aux_schema::block_weight_key(hash)) {
 			match client.header_metadata(hash) {
 				Ok(meta) => {
-					// This should never happen and must be considered a bug.
-					if meta.number <= height_limit {
-						warn!(target: "babe", "unexpected canonical chain state or malformed finality notification");
+					hash = meta.parent;
+
+					// If the parent is part of the canonical chain or there doesn't exist a block
+					// hash for the parent number (bug?!), we can abort adding blocks.
+					if client
+						.hash(meta.number.saturating_sub(1u32.into()))
+						.ok()
+						.flatten()
+						.map_or(true, |h| h == hash)
+					{
 						break
 					}
-					hash = meta.parent;
 				},
 				Err(err) => {
-					warn!(target: "babe", "header lookup fail while cleaning data for block {}: {}", head.to_string(), err.to_string());
+					warn!(
+						target: "babe",
+						"Header lookup fail while cleaning data for block {:?}: {}",
+						hash,
+						err,
+					);
 					break
 				},
 			}

substrate/client/consensus/babe/src/tests.rs

@@ -1043,4 +1043,13 @@ fn obsolete_blocks_aux_data_cleanup() {
 	assert!(aux_data_check(&fork2_hashes, false));
 	// Present C4, C5
 	assert!(aux_data_check(&fork3_hashes, true));
+
+	client.finalize_block(BlockId::Number(4), None, true).unwrap();
+
+	// Wiped: A3
+	assert!(aux_data_check(&fork1_hashes[2..3], false));
+	// Present: A4
+	assert!(aux_data_check(&fork1_hashes[3..], true));
+	// Present C4, C5
+	assert!(aux_data_check(&fork3_hashes, true));
 }