Mmr client gadget - support pallet reset (#12999)

* Remove unneeded code

* Moving some code

* Support pallet-mmr reset

* Rename update_first_mmr_block

Co-authored-by: Adrian Catangiu <adrian@parity.io>

* Renamings

Co-authored-by: Adrian Catangiu <adrian@parity.io>

substrate/client/merkle-mountain-range/src/aux_schema.rs

@@ -54,8 +54,8 @@ fn load_decode<B: AuxStore, T: Decode>(backend: &B, key: &[u8]) -> ClientResult<
 	}
 }
 
-/// Load or initialize persistent data from backend.
-pub(crate) fn load_persistent<B, BE>(backend: &BE) -> ClientResult<Option<PersistedState<B>>>
+/// Load persistent data from backend.
+pub(crate) fn load_state<B, BE>(backend: &BE) -> ClientResult<Option<PersistedState<B>>>
 where
 	B: Block,
 	BE: AuxStore,
@@ -73,15 +73,36 @@ where
 	Ok(None)
 }
 
+/// Load or initialize persistent data from backend.
+pub(crate) fn load_or_init_state<B, BE>(
+	backend: &BE,
+	default: NumberFor<B>,
+) -> sp_blockchain::Result<NumberFor<B>>
+where
+	B: Block,
+	BE: AuxStore,
+{
+	// Initialize gadget best_canon from AUX DB or from pallet genesis.
+	if let Some(best) = load_state::<B, BE>(backend)? {
+		info!(target: LOG_TARGET, "Loading MMR best canonicalized state from db: {:?}.", best);
+		Ok(best)
+	} else {
+		info!(
+			target: LOG_TARGET,
+			"Loading MMR from pallet genesis on what appears to be the first startup: {:?}.",
+			default
+		);
+		write_current_version(backend)?;
+		write_gadget_state::<B, BE>(backend, &default)?;
+		Ok(default)
+	}
+}
+
 #[cfg(test)]
 pub(crate) mod tests {
 	use super::*;
-	use crate::test_utils::{
-		run_test_with_mmr_gadget_pre_post_using_client, MmrBlock, MockClient, OffchainKeyType,
-	};
+	use crate::test_utils::{run_test_with_mmr_gadget_pre_post_using_client, MmrBlock, MockClient};
 	use parking_lot::Mutex;
-	use sp_core::offchain::{DbExternalities, StorageKind};
-	use sp_mmr_primitives::utils::NodesUtils;
 	use sp_runtime::generic::BlockId;
 	use std::{sync::Arc, time::Duration};
 	use substrate_test_runtime_client::{runtime::Block, Backend};
@@ -92,7 +113,7 @@ pub(crate) mod tests {
 		let backend = &*client.backend;
 
 		// version not available in db -> None
-		assert_eq!(load_persistent::<Block, Backend>(backend).unwrap(), None);
+		assert_eq!(load_state::<Block, Backend>(backend).unwrap(), None);
 
 		// populate version in db
 		write_current_version(backend).unwrap();
@@ -100,7 +121,7 @@ pub(crate) mod tests {
 		assert_eq!(load_decode(backend, VERSION_KEY).unwrap(), Some(CURRENT_VERSION));
 
 		// version is available in db but state isn't -> None
-		assert_eq!(load_persistent::<Block, Backend>(backend).unwrap(), None);
+		assert_eq!(load_state::<Block, Backend>(backend).unwrap(), None);
 	}
 
 	#[test]
@@ -113,7 +134,7 @@ pub(crate) mod tests {
 		// version not available in db -> None
 		assert_eq!(load_decode::<Backend, Option<u32>>(&*backend, VERSION_KEY).unwrap(), None);
 		// state not available in db -> None
-		assert_eq!(load_persistent::<Block, Backend>(&*backend).unwrap(), None);
+		assert_eq!(load_state::<Block, Backend>(&*backend).unwrap(), None);
 		// run the gadget while importing and finalizing 3 blocks
 		run_test_with_mmr_gadget_pre_post_using_client(
 			client.clone(),
@@ -136,7 +157,7 @@ pub(crate) mod tests {
 				let backend = &*client.backend;
 				// check there is both version and best canon available in db before running gadget
 				assert_eq!(load_decode(backend, VERSION_KEY).unwrap(), Some(CURRENT_VERSION));
-				assert_eq!(load_persistent::<Block, Backend>(backend).unwrap(), Some(3));
+				assert_eq!(load_state::<Block, Backend>(backend).unwrap(), Some(3));
 			},
 			|client| async move {
 				let a4 = client.import_block(&BlockId::Number(3), b"a4", Some(3)).await;
@@ -148,7 +169,7 @@ pub(crate) mod tests {
 				// a4, a5, a6 were canonicalized
 				client.assert_canonicalized(&[&a4, &a5, &a6]);
 				// check persisted best canon was updated
-				assert_eq!(load_persistent::<Block, Backend>(&*client.backend).unwrap(), Some(6));
+				assert_eq!(load_state::<Block, Backend>(&*client.backend).unwrap(), Some(6));
 			},
 		);
 	}
@@ -177,19 +198,8 @@ pub(crate) mod tests {
 				client.assert_canonicalized(&slice);
 
 				// now manually move them back to non-canon/temp location
-				let mut offchain_db = client.offchain_db();
 				for mmr_block in slice {
-					for node in NodesUtils::right_branch_ending_in_leaf(mmr_block.leaf_idx.unwrap())
-					{
-						let canon_key = mmr_block.get_offchain_key(node, OffchainKeyType::Canon);
-						let val = offchain_db
-							.local_storage_get(StorageKind::PERSISTENT, &canon_key)
-							.unwrap();
-						offchain_db.local_storage_clear(StorageKind::PERSISTENT, &canon_key);
-
-						let temp_key = mmr_block.get_offchain_key(node, OffchainKeyType::Temp);
-						offchain_db.local_storage_set(StorageKind::PERSISTENT, &temp_key, &val);
-					}
+					client.undo_block_canonicalization(mmr_block)
 				}
 			},
 		);
@@ -203,7 +213,7 @@ pub(crate) mod tests {
 				let slice: Vec<&MmrBlock> = blocks.iter().collect();
 
 				// verify persisted state says a1, a2, a3 were canonicalized,
-				assert_eq!(load_persistent::<Block, Backend>(&*client.backend).unwrap(), Some(3));
+				assert_eq!(load_state::<Block, Backend>(&*client.backend).unwrap(), Some(3));
 				// but actually they are NOT canon (we manually reverted them earlier).
 				client.assert_not_canonicalized(&slice);
 			},
@@ -221,7 +231,7 @@ pub(crate) mod tests {
 				// but a4, a5, a6 were canonicalized
 				client.assert_canonicalized(&[&a4, &a5, &a6]);
 				// check persisted best canon was updated
-				assert_eq!(load_persistent::<Block, Backend>(&*client.backend).unwrap(), Some(6));
+				assert_eq!(load_state::<Block, Backend>(&*client.backend).unwrap(), Some(6));
 			},
 		);
 	}

substrate/client/merkle-mountain-range/src/lib.rs

@@ -46,7 +46,7 @@ use crate::offchain_mmr::OffchainMmr;
 use beefy_primitives::MmrRootHash;
 use futures::StreamExt;
 use log::{debug, error, trace, warn};
-use sc_client_api::{Backend, BlockchainEvents, FinalityNotifications};
+use sc_client_api::{Backend, BlockchainEvents, FinalityNotification, FinalityNotifications};
 use sc_offchain::OffchainDb;
 use sp_api::ProvideRuntimeApi;
 use sp_blockchain::{HeaderBackend, HeaderMetadata};
@@ -60,6 +60,62 @@ use std::{marker::PhantomData, sync::Arc};
 /// Logging target for the mmr gadget.
 pub const LOG_TARGET: &str = "mmr";
 
+/// A convenience MMR client trait that defines all the type bounds a MMR client
+/// has to satisfy and defines some helper methods.
+pub trait MmrClient<B, BE>:
+	BlockchainEvents<B> + HeaderBackend<B> + HeaderMetadata<B> + ProvideRuntimeApi<B>
+where
+	B: Block,
+	BE: Backend<B>,
+	Self::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
+{
+	/// Get the block number where the mmr pallet was added to the runtime.
+	fn first_mmr_block_num(&self, notification: &FinalityNotification<B>) -> Option<NumberFor<B>> {
+		let best_block = *notification.header.number();
+		match self.runtime_api().mmr_leaf_count(&BlockId::number(best_block)) {
+			Ok(Ok(mmr_leaf_count)) => {
+				match utils::first_mmr_block_num::<B::Header>(best_block, mmr_leaf_count) {
+					Ok(first_mmr_block) => {
+						debug!(
+							target: LOG_TARGET,
+							"pallet-mmr detected at block {:?} with genesis at block {:?}",
+							best_block,
+							first_mmr_block
+						);
+						Some(first_mmr_block)
+					},
+					Err(e) => {
+						error!(
+							target: LOG_TARGET,
+							"Error calculating the first mmr block: {:?}", e
+						);
+						None
+					},
+				}
+			},
+			_ => {
+				trace!(
+					target: LOG_TARGET,
+					"pallet-mmr not detected at block {:?} ... (best finalized {:?})",
+					best_block,
+					notification.header.number()
+				);
+				None
+			},
+		}
+	}
+}
+
+impl<B, BE, T> MmrClient<B, BE> for T
+where
+	B: Block,
+	BE: Backend<B>,
+	T: BlockchainEvents<B> + HeaderBackend<B> + HeaderMetadata<B> + ProvideRuntimeApi<B>,
+	T::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
+{
+	// empty
+}
+
 struct OffchainMmrBuilder<B: Block, BE: Backend<B>, C> {
 	backend: Arc<BE>,
 	client: Arc<C>,
@@ -73,7 +129,7 @@ impl<B, BE, C> OffchainMmrBuilder<B, BE, C>
 where
 	B: Block,
 	BE: Backend<B>,
-	C: ProvideRuntimeApi<B> + HeaderBackend<B> + HeaderMetadata<B>,
+	C: MmrClient<B, BE>,
 	C::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
 {
 	async fn try_build(
@@ -81,66 +137,21 @@ where
 		finality_notifications: &mut FinalityNotifications<B>,
 	) -> Option<OffchainMmr<B, BE, C>> {
 		while let Some(notification) = finality_notifications.next().await {
-			let best_block = *notification.header.number();
-			match self.client.runtime_api().mmr_leaf_count(&BlockId::number(best_block)) {
-				Ok(Ok(mmr_leaf_count)) => {
-					debug!(
-						target: LOG_TARGET,
-						"pallet-mmr detected at block {:?} with mmr size {:?}",
-						best_block,
-						mmr_leaf_count
-					);
-					match utils::first_mmr_block_num::<B::Header>(best_block, mmr_leaf_count) {
-						Ok(first_mmr_block) => {
-							debug!(
-								target: LOG_TARGET,
-								"pallet-mmr genesis computed at block {:?}", first_mmr_block,
-							);
-							let best_canonicalized =
-								match offchain_mmr::load_or_init_best_canonicalized::<B, BE>(
-									&*self.backend,
-									first_mmr_block,
-								) {
-									Ok(best) => best,
-									Err(e) => {
-										error!(
-											target: LOG_TARGET,
-											"Error loading state from aux db: {:?}", e
-										);
-										return None
-									},
-								};
-							let mut offchain_mmr = OffchainMmr {
-								backend: self.backend,
-								client: self.client,
-								offchain_db: self.offchain_db,
-								indexing_prefix: self.indexing_prefix,
-								first_mmr_block,
-								best_canonicalized,
-							};
-							// We need to make sure all blocks leading up to current notification
-							// have also been canonicalized.
-							offchain_mmr.canonicalize_catch_up(&notification);
-							// We have to canonicalize and prune the blocks in the finality
-							// notification that lead to building the offchain-mmr as well.
-							offchain_mmr.canonicalize_and_prune(notification);
-							return Some(offchain_mmr)
-						},
-						Err(e) => {
-							error!(
-								target: LOG_TARGET,
-								"Error calculating the first mmr block: {:?}", e
-							);
-						},
-					}
-				},
-				_ => {
-					trace!(
-						target: LOG_TARGET,
-						"Waiting for MMR pallet to become available... (best finalized {:?})",
-						notification.header.number()
-					);
-				},
+			if let Some(first_mmr_block_num) = self.client.first_mmr_block_num(&notification) {
+				let mut offchain_mmr = OffchainMmr::new(
+					self.backend,
+					self.client,
+					self.offchain_db,
+					self.indexing_prefix,
+					first_mmr_block_num,
+				)?;
+				// We need to make sure all blocks leading up to current notification
+				// have also been canonicalized.
+				offchain_mmr.canonicalize_catch_up(&notification);
+				// We have to canonicalize and prune the blocks in the finality
+				// notification that lead to building the offchain-mmr as well.
+				offchain_mmr.canonicalize_and_prune(notification);
+				return Some(offchain_mmr)
 			}
 		}
 
@@ -165,7 +176,7 @@ where
 	B: Block,
 	<B::Header as Header>::Number: Into<LeafIndex>,
 	BE: Backend<B>,
-	C: BlockchainEvents<B> + HeaderBackend<B> + HeaderMetadata<B> + ProvideRuntimeApi<B>,
+	C: MmrClient<B, BE>,
 	C::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
 {
 	async fn run(mut self, builder: OffchainMmrBuilder<B, BE, C>) {

substrate/client/merkle-mountain-range/src/offchain_mmr.rs

@@ -21,59 +21,59 @@
 
 #![warn(missing_docs)]
 
-use crate::{aux_schema, LOG_TARGET};
+use crate::{aux_schema, MmrClient, LOG_TARGET};
+use beefy_primitives::MmrRootHash;
 use log::{debug, error, info, warn};
-use sc_client_api::{AuxStore, Backend, FinalityNotification};
+use sc_client_api::{Backend, FinalityNotification};
 use sc_offchain::OffchainDb;
-use sp_blockchain::{CachedHeaderMetadata, ForkBackend, HeaderBackend, HeaderMetadata};
+use sp_blockchain::{CachedHeaderMetadata, ForkBackend};
 use sp_core::offchain::{DbExternalities, StorageKind};
-use sp_mmr_primitives::{utils, utils::NodesUtils, NodeIndex};
+use sp_mmr_primitives::{utils, utils::NodesUtils, MmrApi, NodeIndex};
 use sp_runtime::{
-	traits::{Block, NumberFor, One},
+	traits::{Block, Header, NumberFor, One},
 	Saturating,
 };
 use std::{collections::VecDeque, sync::Arc};
 
-pub(crate) fn load_or_init_best_canonicalized<B, BE>(
-	backend: &BE,
-	first_mmr_block: NumberFor<B>,
-) -> sp_blockchain::Result<NumberFor<B>>
-where
-	BE: AuxStore,
-	B: Block,
-{
-	// Initialize gadget best_canon from AUX DB or from pallet genesis.
-	if let Some(best) = aux_schema::load_persistent::<B, BE>(backend)? {
-		info!(target: LOG_TARGET, "Loading MMR best canonicalized state from db: {:?}.", best);
-		Ok(best)
-	} else {
-		let best = first_mmr_block.saturating_sub(One::one());
-		info!(
-			target: LOG_TARGET,
-			"Loading MMR from pallet genesis on what appears to be the first startup: {:?}.", best
-		);
-		aux_schema::write_current_version(backend)?;
-		aux_schema::write_gadget_state::<B, BE>(backend, &best)?;
-		Ok(best)
-	}
-}
-
 /// `OffchainMMR` exposes MMR offchain canonicalization and pruning logic.
 pub struct OffchainMmr<B: Block, BE: Backend<B>, C> {
-	pub backend: Arc<BE>,
-	pub client: Arc<C>,
-	pub offchain_db: OffchainDb<BE::OffchainStorage>,
-	pub indexing_prefix: Vec<u8>,
-	pub first_mmr_block: NumberFor<B>,
-	pub best_canonicalized: NumberFor<B>,
+	backend: Arc<BE>,
+	client: Arc<C>,
+	offchain_db: OffchainDb<BE::OffchainStorage>,
+	indexing_prefix: Vec<u8>,
+	first_mmr_block: NumberFor<B>,
+	best_canonicalized: NumberFor<B>,
 }
 
 impl<B, BE, C> OffchainMmr<B, BE, C>
 where
-	C: HeaderBackend<B> + HeaderMetadata<B>,
 	BE: Backend<B>,
 	B: Block,
+	C: MmrClient<B, BE>,
+	C::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
 {
+	pub fn new(
+		backend: Arc<BE>,
+		client: Arc<C>,
+		offchain_db: OffchainDb<BE::OffchainStorage>,
+		indexing_prefix: Vec<u8>,
+		first_mmr_block: NumberFor<B>,
+	) -> Option<Self> {
+		let mut best_canonicalized = first_mmr_block.saturating_sub(One::one());
+		best_canonicalized = aux_schema::load_or_init_state::<B, BE>(&*backend, best_canonicalized)
+			.map_err(|e| error!(target: LOG_TARGET, "Error loading state from aux db: {:?}", e))
+			.ok()?;
+
+		Some(Self {
+			backend,
+			client,
+			offchain_db,
+			indexing_prefix,
+			first_mmr_block,
+			best_canonicalized,
+		})
+	}
+
 	fn node_temp_offchain_key(&self, pos: NodeIndex, parent_hash: B::Hash) -> Vec<u8> {
 		NodesUtils::node_temp_offchain_key::<B::Header>(&self.indexing_prefix, pos, parent_hash)
 	}
@@ -82,6 +82,14 @@ where
 		NodesUtils::node_canon_offchain_key(&self.indexing_prefix, pos)
 	}
 
+	fn write_gadget_state_or_log(&self) {
+		if let Err(e) =
+			aux_schema::write_gadget_state::<B, BE>(&*self.backend, &self.best_canonicalized)
+		{
+			debug!(target: LOG_TARGET, "error saving state: {:?}", e);
+		}
+	}
+
 	fn header_metadata_or_log(
 		&self,
 		hash: B::Hash,
@@ -231,10 +239,22 @@ where
 			for hash in to_canon.drain(..) {
 				self.canonicalize_branch(hash);
 			}
-			if let Err(e) =
-				aux_schema::write_gadget_state::<B, BE>(&*self.backend, &self.best_canonicalized)
-			{
-				debug!(target: LOG_TARGET, "error saving state: {:?}", e);
+			self.write_gadget_state_or_log();
+		}
+	}
+
+	fn handle_potential_pallet_reset(&mut self, notification: &FinalityNotification<B>) {
+		if let Some(first_mmr_block_num) = self.client.first_mmr_block_num(&notification) {
+			if first_mmr_block_num != self.first_mmr_block {
+				info!(
+					target: LOG_TARGET,
+					"pallet-mmr reset detected at block {:?} with new genesis at block {:?}",
+					notification.header.number(),
+					first_mmr_block_num
+				);
+				self.first_mmr_block = first_mmr_block_num;
+				self.best_canonicalized = first_mmr_block_num.saturating_sub(One::one());
+				self.write_gadget_state_or_log();
 			}
 		}
 	}
@@ -243,15 +263,14 @@ where
 	/// _canonical key_.
 	/// Prune leafs and nodes added by stale blocks in offchain db from _fork-aware key_.
 	pub fn canonicalize_and_prune(&mut self, notification: FinalityNotification<B>) {
+		// Update the first MMR block in case of a pallet reset.
+		self.handle_potential_pallet_reset(&notification);
+
 		// Move offchain MMR nodes for finalized blocks to canonical keys.
 		for hash in notification.tree_route.iter().chain(std::iter::once(&notification.hash)) {
 			self.canonicalize_branch(*hash);
 		}
-		if let Err(e) =
-			aux_schema::write_gadget_state::<B, BE>(&*self.backend, &self.best_canonicalized)
-		{
-			debug!(target: LOG_TARGET, "error saving state: {:?}", e);
-		}
+		self.write_gadget_state_or_log();
 
 		// Remove offchain MMR nodes for stale forks.
 		let stale_forks = self.client.expand_forks(&notification.stale_heads).unwrap_or_else(
@@ -303,7 +322,7 @@ mod tests {
 			// expected pruned heads because of temp key collision: b1
 			client.assert_pruned(&[&c1, &b1]);
 
-			client.finalize_block(d5.hash(), None);
+			client.finalize_block(d5.hash(), Some(5));
 			tokio::time::sleep(Duration::from_millis(200)).await;
 			// expected finalized heads: d4, d5,
 			client.assert_canonicalized(&[&d4, &d5]);
@@ -312,6 +331,36 @@ mod tests {
 		})
 	}
 
+	#[test]
+	fn canonicalize_and_prune_handles_pallet_reset() {
+		run_test_with_mmr_gadget(|client| async move {
+			// G -> A1 -> A2 -> A3 -> A4 -> A5
+			//      |           |
+			//      |           | -> pallet reset
+			//      |
+			//      | -> first finality notification
+
+			let a1 = client.import_block(&BlockId::Number(0), b"a1", Some(0)).await;
+			let a2 = client.import_block(&BlockId::Hash(a1.hash()), b"a2", Some(1)).await;
+			let a3 = client.import_block(&BlockId::Hash(a2.hash()), b"a3", Some(0)).await;
+			let a4 = client.import_block(&BlockId::Hash(a3.hash()), b"a4", Some(1)).await;
+			let a5 = client.import_block(&BlockId::Hash(a4.hash()), b"a5", Some(2)).await;
+
+			client.finalize_block(a1.hash(), Some(1));
+			tokio::time::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: a1
+			client.assert_canonicalized(&[&a1]);
+			// a2 shouldn't be either canonicalized or pruned. It should be handled as part of the
+			// reset process.
+			client.assert_not_canonicalized(&[&a2]);
+
+			client.finalize_block(a5.hash(), Some(3));
+			tokio::time::sleep(Duration::from_millis(200)).await;
+			//expected finalized heads: a3, a4, a5,
+			client.assert_canonicalized(&[&a3, &a4, &a5]);
+		})
+	}
+
 	#[test]
 	fn canonicalize_catchup_works_correctly() {
 		let mmr_blocks = Arc::new(Mutex::new(vec![]));
@@ -329,11 +378,9 @@ mod tests {
 
 				client.finalize_block(a2.hash(), Some(2));
 
-				{
-					let mut mmr_blocks = mmr_blocks_ref.lock();
-					mmr_blocks.push(a1);
-					mmr_blocks.push(a2);
-				}
+				let mut mmr_blocks = mmr_blocks_ref.lock();
+				mmr_blocks.push(a1);
+				mmr_blocks.push(a2);
 			},
 			|client| async move {
 				// G -> A1 -> A2 -> A3 -> A4
@@ -358,4 +405,54 @@ mod tests {
 			},
 		)
 	}
+
+	#[test]
+	fn canonicalize_catchup_works_correctly_with_pallet_reset() {
+		let mmr_blocks = Arc::new(Mutex::new(vec![]));
+		let mmr_blocks_ref = mmr_blocks.clone();
+		run_test_with_mmr_gadget_pre_post(
+			|client| async move {
+				// G -> A1 -> A2
+				//      |     |
+				//      |     | -> finalized without gadget (missed notification)
+				//      |
+				//      | -> first mmr block
+
+				let a1 = client.import_block(&BlockId::Number(0), b"a1", Some(0)).await;
+				let a2 = client.import_block(&BlockId::Hash(a1.hash()), b"a2", Some(0)).await;
+
+				client.finalize_block(a2.hash(), Some(1));
+
+				let mut mmr_blocks = mmr_blocks_ref.lock();
+				mmr_blocks.push(a1);
+				mmr_blocks.push(a2);
+			},
+			|client| async move {
+				// G -> A1 -> A2 -> A3 -> A4
+				//      |     |     |     |
+				//      |     |     |     | -> finalized after starting gadget
+				//      |     |     |
+				//      |     |     | -> gadget start
+				//      |     |
+				//      |     | -> finalized before gadget start (missed notification)
+				//      |     |    + pallet reset
+				//      |
+				//      | -> first mmr block
+				let blocks = mmr_blocks.lock();
+				let a1 = blocks[0].clone();
+				let a2 = blocks[1].clone();
+				let a3 = client.import_block(&BlockId::Hash(a2.hash()), b"a3", Some(1)).await;
+				let a4 = client.import_block(&BlockId::Hash(a3.hash()), b"a4", Some(2)).await;
+
+				client.finalize_block(a4.hash(), Some(3));
+				tokio::time::sleep(Duration::from_millis(200)).await;
+				// a1 shouldn't be either canonicalized or pruned. It should be handled as part of
+				// the reset process. Checking only that it wasn't pruned. Because of temp key
+				// collision with a2 we can't check that it wasn't canonicalized.
+				client.assert_not_pruned(&[&a1]);
+				// expected finalized heads: a4, a5.
+				client.assert_canonicalized(&[&a2, &a3, &a4]);
+			},
+		)
+	}
 }

substrate/client/merkle-mountain-range/src/test_utils.rs

@@ -162,6 +162,18 @@ impl MockClient {
 		client.finalize_block(hash, None).unwrap();
 	}
 
+	pub fn undo_block_canonicalization(&self, mmr_block: &MmrBlock) {
+		let mut offchain_db = self.offchain_db();
+		for node in NodesUtils::right_branch_ending_in_leaf(mmr_block.leaf_idx.unwrap()) {
+			let canon_key = mmr_block.get_offchain_key(node, OffchainKeyType::Canon);
+			let val = offchain_db.local_storage_get(StorageKind::PERSISTENT, &canon_key).unwrap();
+			offchain_db.local_storage_clear(StorageKind::PERSISTENT, &canon_key);
+
+			let temp_key = mmr_block.get_offchain_key(node, OffchainKeyType::Temp);
+			offchain_db.local_storage_set(StorageKind::PERSISTENT, &temp_key, &val);
+		}
+	}
+
 	pub fn check_offchain_storage<F>(
 		&self,
 		key_type: OffchainKeyType,

substrate/frame/merkle-mountain-range/src/mmr/storage.rs

@@ -60,14 +60,6 @@ impl<StorageType, T, I, L> Default for Storage<StorageType, T, I, L> {
 	}
 }
 
-impl<T, I, L> Storage<OffchainStorage, T, I, L>
-where
-	T: Config<I>,
-	I: 'static,
-	L: primitives::FullLeaf,
-{
-}
-
 impl<T, I, L> mmr_lib::MMRStore<NodeOf<T, I, L>> for Storage<OffchainStorage, T, I, L>
 where
 	T: Config<I>,