From 1185c5b4edd14a7a341329c4dac8b039d5236400 Mon Sep 17 00:00:00 2001
From: varun-doshi <doshivarun202@gmail.com>
Date: Sat, 8 Feb 2025 21:02:56 +0530
Subject: [PATCH] feat: Return root result without blocking due to sparse trie
 Drop

---
 crates/engine/tree/src/tree/root.rs | 36 ++++++++++++++++++++---------
 crates/trie/parallel/src/root.rs    | 14 +++++++++++
 2 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/crates/engine/tree/src/tree/root.rs b/crates/engine/tree/src/tree/root.rs
index 06fbf14ffd4b7..2ac86545d51d2 100644
--- a/crates/engine/tree/src/tree/root.rs
+++ b/crates/engine/tree/src/tree/root.rs
@@ -562,13 +562,7 @@ where
         let (tx, rx) = mpsc::channel();
         thread_pool.spawn(move || {
             debug!(target: "engine::tree", "Starting sparse trie task");
-            let result = match run_sparse_trie(config, metrics, rx) {
-                Ok((state_root, trie_updates, iterations)) => {
-                    StateRootMessage::RootCalculated { state_root, trie_updates, iterations }
-                }
-                Err(error) => StateRootMessage::RootCalculationError(error),
-            };
-            let _ = task_tx.send(result);
+            let _ = run_sparse_trie(config, metrics, rx, task_tx);
         });
         tx
     }
@@ -964,7 +958,8 @@ fn run_sparse_trie<Factory>(
     config: StateRootConfig<Factory>,
     metrics: StateRootTaskMetrics,
     update_rx: mpsc::Receiver<SparseTrieUpdate>,
-) -> Result<(B256, TrieUpdates, u64), ParallelStateRootError>
+    task_tx: Sender<StateRootMessage>,
+) -> Result<(), ParallelStateRootError>
 where
     Factory: DatabaseProviderFactory<Provider: BlockReader> + StateCommitmentProvider,
 {
@@ -1011,12 +1006,31 @@ where
     debug!(target: "engine::root", num_iterations, "All proofs processed, ending calculation");
 
     let start = Instant::now();
-    let root = trie.root().expect("sparse trie should be revealed");
+    let root = trie.root();
     let elapsed = start.elapsed();
     metrics.sparse_trie_final_update_duration_histogram.record(elapsed);
 
-    let trie_updates = trie.take_trie_updates().expect("retention must be enabled");
-    Ok((root, trie_updates, num_iterations))
+    let trie_updates = trie.take_trie_updates();
+
+    let result = match (root, trie_updates) {
+        (Some(state_root), Some(trie_updates)) => StateRootMessage::RootCalculated {
+            state_root,
+            trie_updates,
+            iterations: num_iterations,
+        },
+        (None, None) => StateRootMessage::RootCalculationError(
+            ParallelStateRootError::SparseTreeAndRetentionNotAvailable,
+        ),
+        (None, _) => {
+            StateRootMessage::RootCalculationError(ParallelStateRootError::SparseTreeNotRevealed)
+        }
+        (_, None) => {
+            StateRootMessage::RootCalculationError(ParallelStateRootError::RetentionNotEnabled)
+        }
+    };
+
+    let _ = task_tx.send(result);
+    Ok(())
 }
 
 /// Returns accounts only with those storages that were not already fetched, and
diff --git a/crates/trie/parallel/src/root.rs b/crates/trie/parallel/src/root.rs
index 9ee8ed71e3c45..6aa9079e00858 100644
--- a/crates/trie/parallel/src/root.rs
+++ b/crates/trie/parallel/src/root.rs
@@ -233,6 +233,15 @@ pub enum ParallelStateRootError {
     /// Provider error.
     #[error(transparent)]
     Provider(#[from] ProviderError),
+    /// Sparse Tree not revealed .
+    #[error("sparse tree should be revealed")]
+    SparseTreeNotRevealed,
+    /// Retention not enabled .
+    #[error("retention must be enabled")]
+    RetentionNotEnabled,
+    /// Sparse Tree not revealed and Retention not enabled .
+    #[error("Both root and trie_update should be visible")]
+    SparseTreeAndRetentionNotAvailable,
     /// Other unspecified error.
     #[error("{_0}")]
     Other(String),
@@ -246,6 +255,11 @@ impl From<ParallelStateRootError> for ProviderError {
                 Self::Database(error)
             }
             ParallelStateRootError::Other(other) => Self::Database(DatabaseError::Other(other)),
+            ParallelStateRootError::SparseTreeNotRevealed |
+            ParallelStateRootError::SparseTreeAndRetentionNotAvailable |
+            ParallelStateRootError::RetentionNotEnabled => {
+                Self::Database(DatabaseError::Other(error.to_string()))
+            }
         }
     }
 }