pallet-mmr: move offchain logic to client-side gadget (#12753)

* Move MMR utils methods from pallet to primitives

Signed-off-by: Serban Iorga <serban@parity.io>

* Add method to MmrApi

* Move forks expanding logic from babe to primitives

* Implement MMR gadget

* Remove prunning logic from the MMR pallet

* Code review changes: 1st iteration

* Replace MaybeCanonEngine with CanonEngineBuilder

* fix mmr_leaves_count() for kitchen sink demo

* Update client/merkle-mountain-range/src/canon_engine.rs

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

* Code review changes: 2nd iteration

* fix INDEXING_PREFIX

* impl review comments

* add documentation and minor rename

Signed-off-by: Serban Iorga <serban@parity.io>
Co-authored-by: Adrian Catangiu <adrian@parity.io>

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

@@ -111,7 +111,8 @@ use sp_api::{ApiExt, ProvideRuntimeApi};
 use sp_application_crypto::AppKey;
 use sp_block_builder::BlockBuilder as BlockBuilderApi;
 use sp_blockchain::{
-	Backend as _, Error as ClientError, HeaderBackend, HeaderMetadata, Result as ClientResult,
+	Backend as _, Error as ClientError, ForkBackend, HeaderBackend, HeaderMetadata,
+	Result as ClientResult,
 };
 use sp_consensus::{
 	BlockOrigin, CacheKeyId, Environment, Error as ConsensusError, Proposer, SelectChain,
@@ -123,7 +124,7 @@ use sp_inherents::{CreateInherentDataProviders, InherentData, InherentDataProvid
 use sp_keystore::{SyncCryptoStore, SyncCryptoStorePtr};
 use sp_runtime::{
 	generic::{BlockId, OpaqueDigestItemId},
-	traits::{Block as BlockT, Header, NumberFor, SaturatedConversion, Saturating, Zero},
+	traits::{Block as BlockT, Header, NumberFor, SaturatedConversion, Zero},
 	DigestItem,
 };
 
@@ -515,12 +516,12 @@ fn aux_storage_cleanup<C: HeaderMetadata<Block> + HeaderBackend<Block>, Block: B
 	client: &C,
 	notification: &FinalityNotification<Block>,
 ) -> AuxDataOperations {
-	let mut aux_keys = HashSet::new();
+	let mut hashes = HashSet::new();
 
 	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));
+			hashes.insert(meta.parent);
 		},
 		Err(err) => warn!(
 			target: "babe",
@@ -531,53 +532,29 @@ fn aux_storage_cleanup<C: HeaderMetadata<Block> + HeaderBackend<Block>, Block: B
 	}
 
 	// Cleans data for finalized block's ancestors
-	aux_keys.extend(
+	hashes.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),
+			.filter(|h| **h != notification.hash),
 	);
 
-	// Cleans data for stale branches.
+	// Cleans data for stale forks.
+	let stale_forks = match client.expand_forks(&notification.stale_heads) {
+		Ok(stale_forks) => stale_forks,
+		Err((stale_forks, e)) => {
+			warn!(target: "babe", "{:?}", e,);
+			stale_forks
+		},
+	};
+	hashes.extend(stale_forks.iter());
 
-	for head in notification.stale_heads.iter() {
-		let mut hash = *head;
-		// 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) => {
-					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
-					}
-				},
-				Err(err) => {
-					warn!(
-						target: "babe",
-						"Header lookup fail while cleaning data for block {:?}: {}",
-						hash,
-						err,
-					);
-					break
-				},
-			}
-		}
-	}
-
-	aux_keys.into_iter().map(|val| (val, None)).collect()
+	hashes
+		.into_iter()
+		.map(|val| (aux_schema::block_weight_key(val), None))
+		.collect()
 }
 
 async fn answer_requests<B: BlockT, C>(

substrate/client/merkle-mountain-range/Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "mmr-gadget"
+version = "4.0.0-dev"
+authors = ["Parity Technologies <admin@parity.io>"]
+edition = "2021"
+license = "GPL-3.0-or-later WITH Classpath-exception-2.0"
+repository = "https://github.com/paritytech/substrate"
+description = "MMR Client gadget for substrate"
+homepage = "https://substrate.io"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+codec = { package = "parity-scale-codec", version = "3.0.0" }
+futures = "0.3"
+log = "0.4"
+beefy-primitives = { version = "4.0.0-dev", path = "../../primitives/beefy" }
+sc-client-api = { version = "4.0.0-dev", path = "../api" }
+sp-api = { version = "4.0.0-dev", path = "../../primitives/api" }
+sp-blockchain = { version = "4.0.0-dev", path = "../../primitives/blockchain" }
+sp-consensus = { version = "0.10.0-dev", path = "../../primitives/consensus/common" }
+sp-core = { version = "7.0.0", path = "../../primitives/core" }
+sp-io = { version = "7.0.0", path = "../../primitives/io" }
+sp-mmr-primitives = { version = "4.0.0-dev", path = "../../primitives/merkle-mountain-range" }
+sc-offchain = { version = "4.0.0-dev", path = "../offchain" }
+sp-runtime = { version = "7.0.0", path = "../../primitives/runtime" }
+
+[dev-dependencies]
+tokio = "1.17.0"
+sc-block-builder = { version = "0.10.0-dev", path = "../block-builder" }
+substrate-test-runtime-client = { version = "2.0.0", path = "../../test-utils/runtime/client" }
+async-std = { version = "1.11.0", default-features = false }

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

@@ -0,0 +1,245 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! # MMR offchain gadget
+//!
+//! The MMR offchain gadget is run alongside `pallet-mmr` to assist it with offchain
+//! canonicalization of finalized MMR leaves and nodes.
+//! The gadget should only be run on nodes that have Indexing API enabled (otherwise
+//! `pallet-mmr` cannot write to offchain and this gadget has nothing to do).
+//!
+//! The runtime `pallet-mmr` creates one new MMR leaf per block and all inner MMR parent nodes
+//! generated by the MMR when adding said leaf. MMR nodes are stored both in:
+//! - on-chain storage - hashes only; not full leaf content;
+//! - off-chain storage - via Indexing API, full leaf content (and all internal nodes as well) is
+//!   saved to the Off-chain DB using a key derived from `parent_hash` and node index in MMR. The
+//!   `parent_hash` is also used within the key to avoid conflicts and overwrites on forks (leaf
+//!   data is only allowed to reference data coming from parent block).
+//!
+//! This gadget is driven by block finality and in responsible for pruning stale forks from
+//! offchain db, and moving finalized forks under a "canonical" key based solely on node `pos`
+//! in the MMR.
+
+#![warn(missing_docs)]
+
+mod offchain_mmr;
+#[cfg(test)]
+pub mod test_utils;
+
+use std::{marker::PhantomData, sync::Arc};
+
+use futures::StreamExt;
+use log::{debug, error, trace, warn};
+
+use sc_client_api::{Backend, BlockchainEvents, FinalityNotifications};
+use sc_offchain::OffchainDb;
+use sp_api::ProvideRuntimeApi;
+use sp_blockchain::{HeaderBackend, HeaderMetadata};
+use sp_mmr_primitives::{utils, LeafIndex, MmrApi};
+use sp_runtime::{
+	generic::BlockId,
+	traits::{Block, Header, NumberFor},
+};
+
+use crate::offchain_mmr::OffchainMMR;
+use beefy_primitives::MmrRootHash;
+use sp_core::offchain::OffchainStorage;
+
+/// Logging target for the mmr gadget.
+pub const LOG_TARGET: &str = "mmr";
+
+struct OffchainMmrBuilder<B: Block, C, S> {
+	client: Arc<C>,
+	offchain_db: OffchainDb<S>,
+	indexing_prefix: Vec<u8>,
+
+	_phantom: PhantomData<B>,
+}
+
+impl<B, C, S> OffchainMmrBuilder<B, C, S>
+where
+	B: Block,
+	C: ProvideRuntimeApi<B> + HeaderBackend<B> + HeaderMetadata<B>,
+	C::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
+	S: OffchainStorage,
+{
+	async fn try_build(
+		self,
+		finality_notifications: &mut FinalityNotifications<B>,
+	) -> Option<OffchainMMR<C, B, S>> {
+		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)) => {
+					match utils::first_mmr_block_num::<B::Header>(best_block, mmr_leaf_count) {
+						Ok(first_mmr_block) => {
+							let mut offchain_mmr = OffchainMMR {
+								client: self.client,
+								offchain_db: self.offchain_db,
+								indexing_prefix: self.indexing_prefix,
+								first_mmr_block,
+
+								_phantom: Default::default(),
+							};
+							// 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, "Finality notification: {:?}", notification);
+					debug!(target: LOG_TARGET, "Waiting for MMR pallet to become available ...");
+				},
+			}
+		}
+
+		warn!(
+			target: LOG_TARGET,
+			"Finality notifications stream closed unexpectedly. \
+			Couldn't build the canonicalization engine",
+		);
+		None
+	}
+}
+
+/// A MMR Gadget.
+pub struct MmrGadget<B: Block, BE: Backend<B>, C> {
+	finality_notifications: FinalityNotifications<B>,
+
+	_phantom: PhantomData<(B, BE, C)>,
+}
+
+impl<B, BE, C> MmrGadget<B, BE, C>
+where
+	B: Block,
+	<B::Header as Header>::Number: Into<LeafIndex>,
+	BE: Backend<B>,
+	C: BlockchainEvents<B> + HeaderBackend<B> + HeaderMetadata<B> + ProvideRuntimeApi<B>,
+	C::Api: MmrApi<B, MmrRootHash, NumberFor<B>>,
+{
+	async fn run(mut self, builder: OffchainMmrBuilder<B, C, BE::OffchainStorage>) {
+		let mut offchain_mmr = match builder.try_build(&mut self.finality_notifications).await {
+			Some(offchain_mmr) => offchain_mmr,
+			None => return,
+		};
+
+		while let Some(notification) = self.finality_notifications.next().await {
+			offchain_mmr.canonicalize_and_prune(&notification);
+		}
+	}
+
+	/// Create and run the MMR gadget.
+	pub async fn start(client: Arc<C>, backend: Arc<BE>, indexing_prefix: Vec<u8>) {
+		let offchain_db = match backend.offchain_storage() {
+			Some(offchain_storage) => OffchainDb::new(offchain_storage),
+			None => {
+				warn!(
+					target: LOG_TARGET,
+					"Can't spawn a MmrGadget for a node without offchain storage."
+				);
+				return
+			},
+		};
+
+		let mmr_gadget = MmrGadget::<B, BE, C> {
+			finality_notifications: client.finality_notification_stream(),
+
+			_phantom: Default::default(),
+		};
+		mmr_gadget
+			.run(OffchainMmrBuilder {
+				client,
+				offchain_db,
+				indexing_prefix,
+				_phantom: Default::default(),
+			})
+			.await
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use crate::test_utils::run_test_with_mmr_gadget;
+	use sp_runtime::generic::BlockId;
+	use std::time::Duration;
+
+	#[test]
+	fn mmr_first_block_is_computed_correctly() {
+		// Check the case where the first block is also the first block with MMR.
+		run_test_with_mmr_gadget(|client| async move {
+			// G -> A1 -> A2
+			//      |
+			//      | -> 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(1)).await;
+
+			client.finalize_block(a1.hash(), Some(1));
+			async_std::task::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: a1
+			client.assert_canonicalized(&[&a1]);
+			client.assert_not_pruned(&[&a2]);
+		});
+
+		// Check the case where the first block with MMR comes later.
+		run_test_with_mmr_gadget(|client| async move {
+			// G -> A1 -> A2 -> A3 -> A4 -> A5 -> A6
+			//                        |
+			//                        | -> first mmr block
+
+			let a1 = client.import_block(&BlockId::Number(0), b"a1", None).await;
+			let a2 = client.import_block(&BlockId::Hash(a1.hash()), b"a2", None).await;
+			let a3 = client.import_block(&BlockId::Hash(a2.hash()), b"a3", None).await;
+			let a4 = client.import_block(&BlockId::Hash(a3.hash()), b"a4", Some(0)).await;
+			let a5 = client.import_block(&BlockId::Hash(a4.hash()), b"a5", Some(1)).await;
+			let a6 = client.import_block(&BlockId::Hash(a5.hash()), b"a6", Some(2)).await;
+
+			client.finalize_block(a5.hash(), Some(2));
+			async_std::task::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: a4, a5
+			client.assert_canonicalized(&[&a4, &a5]);
+			client.assert_not_pruned(&[&a6]);
+		});
+	}
+
+	#[test]
+	fn does_not_panic_on_invalid_num_mmr_blocks() {
+		run_test_with_mmr_gadget(|client| async move {
+			// G -> A1
+			//      |
+			//      | -> first mmr block
+
+			let a1 = client.import_block(&BlockId::Number(0), b"a1", Some(0)).await;
+
+			// Simulate the case where the runtime says that there are 2 mmr_blocks when in fact
+			// there is only 1.
+			client.finalize_block(a1.hash(), Some(2));
+			async_std::task::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: -
+			client.assert_not_canonicalized(&[&a1]);
+		});
+	}
+}

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

@@ -0,0 +1,246 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+//! Logic for canonicalizing MMR offchain entries for finalized forks,
+//! and for pruning MMR offchain entries for stale forks.
+
+#![warn(missing_docs)]
+
+use std::{marker::PhantomData, sync::Arc};
+
+use log::{debug, error, warn};
+
+use sc_client_api::FinalityNotification;
+use sc_offchain::OffchainDb;
+use sp_blockchain::{CachedHeaderMetadata, ForkBackend, HeaderBackend, HeaderMetadata};
+use sp_core::offchain::{DbExternalities, OffchainStorage, StorageKind};
+use sp_mmr_primitives::{utils, utils::NodesUtils, NodeIndex};
+use sp_runtime::traits::{Block, Header};
+
+use crate::LOG_TARGET;
+
+/// `OffchainMMR` exposes MMR offchain canonicalization and pruning logic.
+pub struct OffchainMMR<C, B: Block, S> {
+	pub client: Arc<C>,
+	pub offchain_db: OffchainDb<S>,
+	pub indexing_prefix: Vec<u8>,
+	pub first_mmr_block: <B::Header as Header>::Number,
+
+	pub _phantom: PhantomData<B>,
+}
+
+impl<C, S, B> OffchainMMR<C, B, S>
+where
+	C: HeaderBackend<B> + HeaderMetadata<B>,
+	S: OffchainStorage,
+	B: Block,
+{
+	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)
+	}
+
+	fn node_canon_offchain_key(&self, pos: NodeIndex) -> Vec<u8> {
+		NodesUtils::node_canon_offchain_key(&self.indexing_prefix, pos)
+	}
+
+	fn header_metadata_or_log(
+		&self,
+		hash: B::Hash,
+		action: &str,
+	) -> Option<CachedHeaderMetadata<B>> {
+		match self.client.header_metadata(hash) {
+			Ok(header) => Some(header),
+			_ => {
+				error!(
+					target: LOG_TARGET,
+					"Block {} not found. Couldn't {} associated branch.", hash, action
+				);
+				None
+			},
+		}
+	}
+
+	fn right_branch_ending_in_block_or_log(
+		&self,
+		block_num: <B::Header as Header>::Number,
+		action: &str,
+	) -> Option<Vec<NodeIndex>> {
+		match utils::block_num_to_leaf_index::<B::Header>(block_num, self.first_mmr_block) {
+			Ok(leaf_idx) => {
+				let branch = NodesUtils::right_branch_ending_in_leaf(leaf_idx);
+				debug!(
+					target: LOG_TARGET,
+					"Nodes to {} for block {}: {:?}", action, block_num, branch
+				);
+				Some(branch)
+			},
+			Err(e) => {
+				error!(
+					target: LOG_TARGET,
+					"Error converting block number {} to leaf index: {:?}. \
+					Couldn't {} associated branch.",
+					block_num,
+					e,
+					action
+				);
+				None
+			},
+		}
+	}
+
+	fn prune_branch(&mut self, block_hash: &B::Hash) {
+		let action = "prune";
+		let header = match self.header_metadata_or_log(*block_hash, action) {
+			Some(header) => header,
+			_ => return,
+		};
+
+		// We prune the leaf associated with the provided block and all the nodes added by that
+		// leaf.
+		let stale_nodes = match self.right_branch_ending_in_block_or_log(header.number, action) {
+			Some(nodes) => nodes,
+			None => {
+				// If we can't convert the block number to a leaf index, the chain state is probably
+				// corrupted. We only log the error, hoping that the chain state will be fixed.
+				return
+			},
+		};
+
+		for pos in stale_nodes {
+			let temp_key = self.node_temp_offchain_key(pos, header.parent);
+			self.offchain_db.local_storage_clear(StorageKind::PERSISTENT, &temp_key);
+			debug!(target: LOG_TARGET, "Pruned elem at pos {} with temp key {:?}", pos, temp_key);
+		}
+	}
+
+	fn canonicalize_branch(&mut self, block_hash: &B::Hash) {
+		let action = "canonicalize";
+		let header = match self.header_metadata_or_log(*block_hash, action) {
+			Some(header) => header,
+			_ => return,
+		};
+
+		// Don't canonicalize branches corresponding to blocks for which the MMR pallet
+		// wasn't yet initialized.
+		if header.number < self.first_mmr_block {
+			return
+		}
+
+		// We "canonicalize" the leaf associated with the provided block
+		// and all the nodes added by that leaf.
+		let to_canon_nodes = match self.right_branch_ending_in_block_or_log(header.number, action) {
+			Some(nodes) => nodes,
+			None => {
+				// If we can't convert the block number to a leaf index, the chain state is probably
+				// corrupted. We only log the error, hoping that the chain state will be fixed.
+				return
+			},
+		};
+
+		for pos in to_canon_nodes {
+			let temp_key = self.node_temp_offchain_key(pos, header.parent);
+			if let Some(elem) =
+				self.offchain_db.local_storage_get(StorageKind::PERSISTENT, &temp_key)
+			{
+				let canon_key = self.node_canon_offchain_key(pos);
+				self.offchain_db.local_storage_set(StorageKind::PERSISTENT, &canon_key, &elem);
+				self.offchain_db.local_storage_clear(StorageKind::PERSISTENT, &temp_key);
+				debug!(
+					target: LOG_TARGET,
+					"Moved elem at pos {} from temp key {:?} to canon key {:?}",
+					pos,
+					temp_key,
+					canon_key
+				);
+			} else {
+				error!(
+					target: LOG_TARGET,
+					"Couldn't canonicalize elem at pos {} using temp key {:?}", pos, temp_key
+				);
+			}
+		}
+	}
+
+	/// Move leafs and nodes added by finalized blocks in offchain db from _fork-aware key_ to
+	/// _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>) {
+		// Move offchain MMR nodes for finalized blocks to canonical keys.
+		for block_hash in notification.tree_route.iter().chain(std::iter::once(&notification.hash))
+		{
+			self.canonicalize_branch(block_hash);
+		}
+
+		// Remove offchain MMR nodes for stale forks.
+		let stale_forks = self.client.expand_forks(&notification.stale_heads).unwrap_or_else(
+			|(stale_forks, e)| {
+				warn!(target: LOG_TARGET, "{:?}", e);
+				stale_forks
+			},
+		);
+		for hash in stale_forks.iter() {
+			self.prune_branch(hash);
+		}
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use crate::test_utils::run_test_with_mmr_gadget;
+	use sp_runtime::generic::BlockId;
+	use std::time::Duration;
+
+	#[test]
+	fn canonicalize_and_prune_works_correctly() {
+		run_test_with_mmr_gadget(|client| async move {
+			//                     -> D4 -> D5
+			// G -> A1 -> A2 -> A3 -> A4
+			//   -> B1 -> B2 -> B3
+			//   -> C1
+
+			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(2)).await;
+			let a4 = client.import_block(&BlockId::Hash(a3.hash()), b"a4", Some(3)).await;
+
+			let b1 = client.import_block(&BlockId::Number(0), b"b1", Some(0)).await;
+			let b2 = client.import_block(&BlockId::Hash(b1.hash()), b"b2", Some(1)).await;
+			let b3 = client.import_block(&BlockId::Hash(b2.hash()), b"b3", Some(2)).await;
+
+			let c1 = client.import_block(&BlockId::Number(0), b"c1", Some(0)).await;
+
+			let d4 = client.import_block(&BlockId::Hash(a3.hash()), b"d4", Some(3)).await;
+			let d5 = client.import_block(&BlockId::Hash(d4.hash()), b"d5", Some(4)).await;
+
+			client.finalize_block(a3.hash(), Some(3));
+			async_std::task::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: a1, a2, a3
+			client.assert_canonicalized(&[&a1, &a2, &a3]);
+			// expected stale heads: c1
+			// expected pruned heads because of temp key collision: b1
+			client.assert_pruned(&[&c1, &b1]);
+
+			client.finalize_block(d5.hash(), None);
+			async_std::task::sleep(Duration::from_millis(200)).await;
+			// expected finalized heads: d4, d5,
+			client.assert_canonicalized(&[&d4, &d5]);
+			// expected stale heads: b1, b2, b3, a4
+			client.assert_pruned(&[&b1, &b2, &b3, &a4]);
+		})
+	}
+}

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

@@ -0,0 +1,344 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2021-2022 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use futures::{executor::LocalPool, task::LocalSpawn, FutureExt};
+use std::{
+	future::Future,
+	sync::{Arc, Mutex},
+	time::Duration,
+};
+
+use sc_block_builder::BlockBuilderProvider;
+use sc_client_api::{
+	Backend as BackendT, BlockchainEvents, FinalityNotifications, ImportNotifications,
+	StorageEventStream, StorageKey,
+};
+use sc_offchain::OffchainDb;
+use sp_api::{ApiRef, ProvideRuntimeApi};
+use sp_blockchain::{BlockStatus, CachedHeaderMetadata, HeaderBackend, HeaderMetadata, Info};
+use sp_consensus::BlockOrigin;
+use sp_core::{
+	offchain::{DbExternalities, StorageKind},
+	H256,
+};
+use sp_mmr_primitives as mmr;
+use sp_mmr_primitives::{utils::NodesUtils, LeafIndex, NodeIndex};
+use sp_runtime::{
+	generic::BlockId,
+	traits::{Block as BlockT, Header as HeaderT},
+};
+use substrate_test_runtime_client::{
+	runtime::{Block, BlockNumber, Hash, Header},
+	Backend, BlockBuilderExt, Client, ClientBlockImportExt, ClientExt, DefaultTestClientBuilderExt,
+	TestClientBuilder, TestClientBuilderExt,
+};
+
+use crate::MmrGadget;
+
+type MmrHash = H256;
+
+struct MockRuntimeApiData {
+	num_blocks: BlockNumber,
+}
+
+#[derive(Clone)]
+pub struct MockRuntimeApi {
+	data: Arc<Mutex<MockRuntimeApiData>>,
+}
+
+impl MockRuntimeApi {
+	pub const INDEXING_PREFIX: &'static [u8] = b"mmr_test";
+}
+
+pub struct MmrBlock {
+	block: Block,
+	leaf_idx: Option<LeafIndex>,
+	leaf_data: Vec<u8>,
+}
+
+#[derive(Clone, Copy)]
+pub enum OffchainKeyType {
+	Temp,
+	Canon,
+}
+
+impl MmrBlock {
+	pub fn hash(&self) -> Hash {
+		self.block.hash()
+	}
+
+	pub fn parent_hash(&self) -> Hash {
+		*self.block.header.parent_hash()
+	}
+
+	pub fn get_offchain_key(&self, node: NodeIndex, key_type: OffchainKeyType) -> Vec<u8> {
+		match key_type {
+			OffchainKeyType::Temp => NodesUtils::node_temp_offchain_key::<Header>(
+				MockRuntimeApi::INDEXING_PREFIX,
+				node,
+				*self.block.header.parent_hash(),
+			),
+			OffchainKeyType::Canon =>
+				NodesUtils::node_canon_offchain_key(MockRuntimeApi::INDEXING_PREFIX, node),
+		}
+	}
+}
+
+pub struct MockClient {
+	client: Mutex<Client<Backend>>,
+	backend: Arc<Backend>,
+	runtime_api_params: Arc<Mutex<MockRuntimeApiData>>,
+}
+
+impl MockClient {
+	fn new() -> Self {
+		let client_builder = TestClientBuilder::new().enable_offchain_indexing_api();
+		let (client, backend) = client_builder.build_with_backend();
+		MockClient {
+			client: Mutex::new(client),
+			backend,
+			runtime_api_params: Arc::new(Mutex::new(MockRuntimeApiData { num_blocks: 0 })),
+		}
+	}
+
+	fn offchain_db(&self) -> OffchainDb<<Backend as BackendT<Block>>::OffchainStorage> {
+		OffchainDb::new(self.backend.offchain_storage().unwrap())
+	}
+
+	pub async fn import_block(
+		&self,
+		at: &BlockId<Block>,
+		name: &[u8],
+		maybe_leaf_idx: Option<LeafIndex>,
+	) -> MmrBlock {
+		let mut client = self.client.lock().unwrap();
+
+		let mut block_builder = client.new_block_at(at, Default::default(), false).unwrap();
+		// Make sure the block has a different hash than its siblings
+		block_builder
+			.push_storage_change(b"name".to_vec(), Some(name.to_vec()))
+			.unwrap();
+		let block = block_builder.build().unwrap().block;
+		client.import(BlockOrigin::Own, block.clone()).await.unwrap();
+
+		let parent_hash = *block.header.parent_hash();
+		// Simulate writing MMR nodes in offchain storage
+		if let Some(leaf_idx) = maybe_leaf_idx {
+			let mut offchain_db = self.offchain_db();
+			for node in NodesUtils::right_branch_ending_in_leaf(leaf_idx) {
+				let temp_key = NodesUtils::node_temp_offchain_key::<Header>(
+					MockRuntimeApi::INDEXING_PREFIX,
+					node,
+					parent_hash,
+				);
+				offchain_db.local_storage_set(
+					StorageKind::PERSISTENT,
+					&temp_key,
+					parent_hash.as_ref(),
+				)
+			}
+		}
+
+		MmrBlock { block, leaf_idx: maybe_leaf_idx, leaf_data: parent_hash.as_ref().to_vec() }
+	}
+
+	pub fn finalize_block(&self, hash: Hash, maybe_num_mmr_blocks: Option<BlockNumber>) {
+		let client = self.client.lock().unwrap();
+		if let Some(num_mmr_blocks) = maybe_num_mmr_blocks {
+			self.runtime_api_params.lock().unwrap().num_blocks = num_mmr_blocks;
+		}
+
+		client.finalize_block(hash, None).unwrap();
+	}
+
+	pub fn check_offchain_storage<F>(
+		&self,
+		key_type: OffchainKeyType,
+		blocks: &[&MmrBlock],
+		mut f: F,
+	) where
+		F: FnMut(Option<Vec<u8>>, &MmrBlock),
+	{
+		let mut offchain_db = self.offchain_db();
+		for mmr_block in blocks {
+			for node in NodesUtils::right_branch_ending_in_leaf(mmr_block.leaf_idx.unwrap()) {
+				let temp_key = mmr_block.get_offchain_key(node, key_type);
+				let val = offchain_db.local_storage_get(StorageKind::PERSISTENT, &temp_key);
+				f(val, mmr_block);
+			}
+		}
+	}
+
+	pub fn assert_pruned(&self, blocks: &[&MmrBlock]) {
+		self.check_offchain_storage(OffchainKeyType::Temp, blocks, |val, _block| {
+			assert!(val.is_none());
+		})
+	}
+
+	pub fn assert_not_pruned(&self, blocks: &[&MmrBlock]) {
+		self.check_offchain_storage(OffchainKeyType::Temp, blocks, |val, block| {
+			assert_eq!(val.as_ref(), Some(&block.leaf_data));
+		})
+	}
+
+	pub fn assert_canonicalized(&self, blocks: &[&MmrBlock]) {
+		self.check_offchain_storage(OffchainKeyType::Canon, blocks, |val, block| {
+			assert_eq!(val.as_ref(), Some(&block.leaf_data));
+		});
+
+		self.assert_pruned(blocks);
+	}
+
+	pub fn assert_not_canonicalized(&self, blocks: &[&MmrBlock]) {
+		self.check_offchain_storage(OffchainKeyType::Canon, blocks, |val, _block| {
+			assert!(val.is_none());
+		});
+
+		self.assert_not_pruned(blocks);
+	}
+}
+
+impl HeaderMetadata<Block> for MockClient {
+	type Error = <Client<Backend> as HeaderMetadata<Block>>::Error;
+
+	fn header_metadata(&self, hash: Hash) -> Result<CachedHeaderMetadata<Block>, Self::Error> {
+		self.client.lock().unwrap().header_metadata(hash)
+	}
+
+	fn insert_header_metadata(&self, _hash: Hash, _header_metadata: CachedHeaderMetadata<Block>) {
+		todo!()
+	}
+
+	fn remove_header_metadata(&self, _hash: Hash) {
+		todo!()
+	}
+}
+
+impl HeaderBackend<Block> for MockClient {
+	fn header(&self, id: BlockId<Block>) -> sc_client_api::blockchain::Result<Option<Header>> {
+		self.client.lock().unwrap().header(&id)
+	}
+
+	fn info(&self) -> Info<Block> {
+		self.client.lock().unwrap().info()
+	}
+
+	fn status(&self, id: BlockId<Block>) -> sc_client_api::blockchain::Result<BlockStatus> {
+		self.client.lock().unwrap().status(id)
+	}
+
+	fn number(&self, hash: Hash) -> sc_client_api::blockchain::Result<Option<BlockNumber>> {
+		self.client.lock().unwrap().number(hash)
+	}
+
+	fn hash(&self, number: BlockNumber) -> sc_client_api::blockchain::Result<Option<Hash>> {
+		self.client.lock().unwrap().hash(number)
+	}
+}
+
+impl BlockchainEvents<Block> for MockClient {
+	fn import_notification_stream(&self) -> ImportNotifications<Block> {
+		unimplemented!()
+	}
+
+	fn finality_notification_stream(&self) -> FinalityNotifications<Block> {
+		self.client.lock().unwrap().finality_notification_stream()
+	}
+
+	fn storage_changes_notification_stream(
+		&self,
+		_filter_keys: Option<&[StorageKey]>,
+		_child_filter_keys: Option<&[(StorageKey, Option<Vec<StorageKey>>)]>,
+	) -> sc_client_api::blockchain::Result<StorageEventStream<Hash>> {
+		unimplemented!()
+	}
+}
+
+impl ProvideRuntimeApi<Block> for MockClient {
+	type Api = MockRuntimeApi;
+
+	fn runtime_api(&self) -> ApiRef<'_, Self::Api> {
+		MockRuntimeApi { data: self.runtime_api_params.clone() }.into()
+	}
+}
+
+sp_api::mock_impl_runtime_apis! {
+	impl mmr::MmrApi<Block, MmrHash, BlockNumber> for MockRuntimeApi {
+		fn mmr_root() -> Result<MmrHash, mmr::Error> {
+			Err(mmr::Error::PalletNotIncluded)
+		}
+
+		fn mmr_leaf_count(&self) -> Result<LeafIndex, mmr::Error> {
+			Ok(self.data.lock().unwrap().num_blocks)
+		}
+
+		fn generate_proof(
+			&self,
+			_block_numbers: Vec<u64>,
+			_best_known_block_number: Option<u64>,
+		) -> Result<(Vec<mmr::EncodableOpaqueLeaf>, mmr::Proof<MmrHash>), mmr::Error> {
+			Err(mmr::Error::PalletNotIncluded)
+		}
+
+		fn verify_proof(_leaves: Vec<mmr::EncodableOpaqueLeaf>, _proof: mmr::Proof<MmrHash>)
+			-> Result<(), mmr::Error>
+		{
+			Err(mmr::Error::PalletNotIncluded)
+		}
+
+		fn verify_proof_stateless(
+			_root: MmrHash,
+			_leaves: Vec<mmr::EncodableOpaqueLeaf>,
+			_proof: mmr::Proof<MmrHash>
+		) -> Result<(), mmr::Error> {
+			Err(mmr::Error::PalletNotIncluded)
+		}
+	}
+}
+
+pub fn run_test_with_mmr_gadget<F, Fut>(f: F)
+where
+	F: FnOnce(Arc<MockClient>) -> Fut + 'static,
+	Fut: Future<Output = ()>,
+{
+	let mut pool = LocalPool::new();
+	let client = Arc::new(MockClient::new());
+
+	let client_clone = client.clone();
+	pool.spawner()
+		.spawn_local_obj(
+			async move {
+				let backend = client_clone.backend.clone();
+				MmrGadget::start(
+					client_clone.clone(),
+					backend,
+					MockRuntimeApi::INDEXING_PREFIX.to_vec(),
+				)
+				.await
+			}
+			.boxed_local()
+			.into(),
+		)
+		.unwrap();
+
+	pool.run_until(async move {
+		async_std::task::sleep(Duration::from_millis(200)).await;
+
+		f(client).await
+	});
+}