Mmr persist state (#12822)

client/mmr: persisting gadget state across runs

Fixes #12780

* client/mmr: on init do canonicalization catch-up

* client/mmr: add more tests

* client/mmr: persist gadget progress in aux db

* client/mmr: add more tests

* client/mmr: replace async_std with tokio

* remove leftover comment

* address review comments

Signed-off-by: acatangiu <adrian@parity.io>

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

@@ -37,15 +37,15 @@
 
 #![warn(missing_docs)]
 
+mod aux_schema;
 mod offchain_mmr;
 #[cfg(test)]
 pub mod test_utils;
 
-use std::{marker::PhantomData, sync::Arc};
-
+use crate::offchain_mmr::OffchainMmr;
+use beefy_primitives::MmrRootHash;
 use futures::StreamExt;
-use log::{error, trace, warn};
-
+use log::{debug, error, trace, warn};
 use sc_client_api::{Backend, BlockchainEvents, FinalityNotifications};
 use sc_offchain::OffchainDb;
 use sp_api::ProvideRuntimeApi;
@@ -55,50 +55,75 @@ use sp_runtime::{
 	generic::BlockId,
 	traits::{Block, Header, NumberFor},
 };
-
-use crate::offchain_mmr::OffchainMMR;
-use beefy_primitives::MmrRootHash;
-use sp_core::offchain::OffchainStorage;
+use std::{marker::PhantomData, sync::Arc};
 
 /// Logging target for the mmr gadget.
 pub const LOG_TARGET: &str = "mmr";
 
-struct OffchainMmrBuilder<B: Block, C, S> {
+struct OffchainMmrBuilder<B: Block, BE: Backend<B>, C> {
+	backend: Arc<BE>,
 	client: Arc<C>,
-	offchain_db: OffchainDb<S>,
+	offchain_db: OffchainDb<BE::OffchainStorage>,
 	indexing_prefix: Vec<u8>,
 
 	_phantom: PhantomData<B>,
 }
 
-impl<B, C, S> OffchainMmrBuilder<B, C, S>
+impl<B, BE, C> OffchainMmrBuilder<B, BE, C>
 where
 	B: Block,
+	BE: Backend<B>,
 	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>> {
+	) -> 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) => {
-							let mut offchain_mmr = OffchainMMR {
+							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,
-
-								_phantom: Default::default(),
+								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);
+							offchain_mmr.canonicalize_and_prune(notification);
 							return Some(offchain_mmr)
 						},
 						Err(e) => {
@@ -143,14 +168,14 @@ where
 	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>) {
+	async fn run(mut self, builder: OffchainMmrBuilder<B, BE, C>) {
 		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);
+			offchain_mmr.canonicalize_and_prune(notification);
 		}
 	}
 
@@ -174,6 +199,7 @@ where
 		};
 		mmr_gadget
 			.run(OffchainMmrBuilder {
+				backend,
 				client,
 				offchain_db,
 				indexing_prefix,