Merge pull request #815 from worldcoin/piohei/fix_tree_restore

Fix tree restore on startup.

src/identity/processor.rs

@@ -18,7 +18,9 @@ use crate::database::methods::DbMethods;
 use crate::database::types::{BatchEntry, BatchType};
 use crate::database::{Database, IsolationLevel};
 use crate::ethereum::{Ethereum, ReadProvider};
-use crate::identity_tree::{Canonical, Hash, Intermediate, TreeVersion, TreeWithNextVersion};
+use crate::identity_tree::{
+    Canonical, Hash, Intermediate, ProcessedStatus, TreeVersion, TreeWithNextVersion,
+};
 use crate::prover::identity::Identity;
 use crate::prover::repository::ProverRepository;
 use crate::prover::Prover;
@@ -41,6 +43,8 @@ pub trait IdentityProcessor: Send + Sync + 'static {
     async fn mine_transaction(&self, transaction_id: TransactionId) -> anyhow::Result<bool>;
 
     async fn tree_init_correction(&self, initial_root_hash: &Hash) -> anyhow::Result<()>;
+
+    async fn latest_root(&self) -> anyhow::Result<Option<Hash>>;
 }
 
 pub struct OnChainIdentityProcessor {
@@ -158,6 +162,10 @@ impl IdentityProcessor for OnChainIdentityProcessor {
 
         Ok(())
     }
+
+    async fn latest_root(&self) -> anyhow::Result<Option<Hash>> {
+        Ok(Some(self.identity_manager.latest_root().await?.into()))
+    }
 }
 
 impl OnChainIdentityProcessor {
@@ -561,6 +569,13 @@ impl IdentityProcessor for OffChainIdentityProcessor {
         // For off chain mode we don't correct tree at all
         Ok(())
     }
+
+    async fn latest_root(&self) -> anyhow::Result<Option<Hash>> {
+        Ok(self
+            .database
+            .get_latest_root_by_status(ProcessedStatus::Mined)
+            .await?)
+    }
 }
 
 impl OffChainIdentityProcessor {

src/identity_tree/initializer.rs

@@ -47,27 +47,39 @@ impl TreeInitializer {
             .await?;
 
         let timer = Instant::now();
-        let mut tree_state = self.restore_or_initialize_tree(initial_root_hash).await?;
+        info!("Tree state initialization started");
+        let tree_state = self
+            .restore_or_initialize_tree(initial_root_hash, self.config.force_cache_purge)
+            .await?;
         info!("Tree state initialization took: {:?}", timer.elapsed());
 
         let tree_root = tree_state.get_processed_tree().get_root();
-
-        if tree_root != initial_root_hash {
-            warn!(
-                "Cached tree root is different from the contract root. Purging cache and \
+        match self.identity_processor.latest_root().await? {
+            Some(root) if root == tree_root => Ok(tree_state),
+            None if initial_root_hash == tree_root => Ok(tree_state),
+            _ => {
+                warn!(
+                    "Cached tree root is different from the contract root. Purging cache and \
                  reinitializing."
-            );
+                );
 
-            tree_state = self.restore_or_initialize_tree(initial_root_hash).await?;
-        }
+                let timer = Instant::now();
+                info!("Tree state initialization started");
+                let tree_state = self
+                    .restore_or_initialize_tree(initial_root_hash, true)
+                    .await?;
+                info!("Tree state initialization took: {:?}", timer.elapsed());
 
-        Ok(tree_state)
+                Ok(tree_state)
+            }
+        }
     }
 
     #[instrument(skip(self))]
     async fn restore_or_initialize_tree(
         &self,
         initial_root_hash: Hash,
+        force_cache_purge: bool,
     ) -> anyhow::Result<TreeState> {
         let mut mined_items = self
             .database
@@ -78,7 +90,7 @@ impl TreeInitializer {
 
         let mined_items = dedup_tree_updates(mined_items);
 
-        if !self.config.force_cache_purge {
+        if !force_cache_purge {
             info!("Attempting to restore tree from cache");
             if let Some(tree_state) = self
                 .get_cached_tree_state(&mined_items, initial_root_hash)