fix(tree): use in memory trie cursor for trie updates comparison in tree (#13789)

crates/engine/tree/src/tree/mod.rs

@@ -2276,52 +2276,60 @@ where
         let persistence_not_in_progress = !self.persistence_state.in_progress();
 
         let state_root_result = std::thread::scope(|scope| {
-            let (state_root_handle, state_hook) = if persistence_not_in_progress &&
-                self.config.use_state_root_task()
-            {
-                let consistent_view = ConsistentDbView::new_with_latest_tip(self.provider.clone())?;
+            let (state_root_handle, in_memory_trie_cursor, state_hook) =
+                if persistence_not_in_progress && self.config.use_state_root_task() {
+                    let consistent_view =
+                        ConsistentDbView::new_with_latest_tip(self.provider.clone())?;
 
-                let state_root_config = StateRootConfig::new_from_input(
-                    consistent_view.clone(),
-                    self.compute_trie_input(consistent_view.clone(), block.header().parent_hash())
+                    let state_root_config = StateRootConfig::new_from_input(
+                        consistent_view.clone(),
+                        self.compute_trie_input(
+                            consistent_view.clone(),
+                            block.header().parent_hash(),
+                        )
                         .map_err(|e| InsertBlockErrorKind::Other(Box::new(e)))?,
-                );
+                    );
 
-                let provider_ro = consistent_view.provider_ro()?;
-                let nodes_sorted = state_root_config.nodes_sorted.clone();
-                let state_sorted = state_root_config.state_sorted.clone();
-                let prefix_sets = state_root_config.prefix_sets.clone();
+                    let provider_ro = consistent_view.provider_ro()?;
+                    let nodes_sorted = state_root_config.nodes_sorted.clone();
+                    let state_sorted = state_root_config.state_sorted.clone();
+                    let prefix_sets = state_root_config.prefix_sets.clone();
 
-                // context will hold the values that need to be kept alive
-                let context =
-                    StateHookContext { provider_ro, nodes_sorted, state_sorted, prefix_sets };
+                    // context will hold the values that need to be kept alive
+                    let context =
+                        StateHookContext { provider_ro, nodes_sorted, state_sorted, prefix_sets };
 
-                // it is ok to leak here because we are in a scoped thread, the
-                // memory will be freed when the thread completes
-                let context = Box::leak(Box::new(context));
+                    // it is ok to leak here because we are in a scoped thread, the
+                    // memory will be freed when the thread completes
+                    let context = Box::leak(Box::new(context));
 
-                let blinded_provider_factory = ProofBlindedProviderFactory::new(
-                    InMemoryTrieCursorFactory::new(
+                    let in_memory_trie_cursor = InMemoryTrieCursorFactory::new(
                         DatabaseTrieCursorFactory::new(context.provider_ro.tx_ref()),
                         &context.nodes_sorted,
-                    ),
-                    HashedPostStateCursorFactory::new(
-                        DatabaseHashedCursorFactory::new(context.provider_ro.tx_ref()),
-                        &context.state_sorted,
-                    ),
-                    context.prefix_sets.clone(),
-                );
+                    );
+                    let blinded_provider_factory = ProofBlindedProviderFactory::new(
+                        in_memory_trie_cursor.clone(),
+                        HashedPostStateCursorFactory::new(
+                            DatabaseHashedCursorFactory::new(context.provider_ro.tx_ref()),
+                            &context.state_sorted,
+                        ),
+                        context.prefix_sets.clone(),
+                    );
 
-                let state_root_task = StateRootTask::new(
-                    state_root_config,
-                    blinded_provider_factory,
-                    self.state_root_task_pool.clone(),
-                );
-                let state_hook = state_root_task.state_hook();
-                (Some(state_root_task.spawn(scope)), Box::new(state_hook) as Box<dyn OnStateHook>)
-            } else {
-                (None, Box::new(|_state: &EvmState| {}) as Box<dyn OnStateHook>)
-            };
+                    let state_root_task = StateRootTask::new(
+                        state_root_config,
+                        blinded_provider_factory,
+                        self.state_root_task_pool.clone(),
+                    );
+                    let state_hook = state_root_task.state_hook();
+                    (
+                        Some(state_root_task.spawn(scope)),
+                        Some(in_memory_trie_cursor),
+                        Box::new(state_hook) as Box<dyn OnStateHook>,
+                    )
+                } else {
+                    (None, None, Box::new(|_state: &EvmState| {}) as Box<dyn OnStateHook>)
+                };
 
             let execution_start = Instant::now();
             let output = self.metrics.executor.execute_metered(executor, &block, state_hook)?;
@@ -2383,9 +2391,8 @@ where
                                     state_provider.state_root_with_updates(hashed_state.clone())?;
 
                                 if regular_root == block.header().state_root() {
-                                    let provider = self.provider.database_provider_ro()?;
                                     compare_trie_updates(
-                                        provider.tx_ref(),
+                                        in_memory_trie_cursor.expect("in memory trie cursor must exist if use_state_root_task is true"),
                                         task_trie_updates.clone(),
                                         regular_updates,
                                     )

crates/engine/tree/src/tree/trie_updates.rs

@@ -1,14 +1,12 @@
-use alloy_primitives::{
-    map::{HashMap, HashSet},
-    B256,
-};
-use reth_db::{transaction::DbTx, DatabaseError};
+use std::collections::BTreeSet;
+
+use alloy_primitives::{map::HashMap, B256};
+use reth_db::DatabaseError;
 use reth_trie::{
     trie_cursor::{TrieCursor, TrieCursorFactory},
     updates::{StorageTrieUpdates, TrieUpdates},
     BranchNodeCompact, Nibbles,
 };
-use reth_trie_db::DatabaseTrieCursorFactory;
 use tracing::debug;
 
 #[derive(Debug)]
@@ -38,8 +36,16 @@ impl TrieUpdatesDiff {
                 debug!(target: "engine::tree", ?path, ?task, ?regular, ?database, "Difference in account trie updates");
             }
 
-            for (path, EntryDiff { task, regular, database }) in &self.removed_nodes {
-                debug!(target: "engine::tree", ?path, ?task, ?regular, ?database, "Difference in removed account trie nodes");
+            for (
+                path,
+                EntryDiff {
+                    task: task_removed,
+                    regular: regular_removed,
+                    database: database_not_exists,
+                },
+            ) in &self.removed_nodes
+            {
+                debug!(target: "engine::tree", ?path, ?task_removed, ?regular_removed, ?database_not_exists, "Difference in removed account trie nodes");
             }
 
             for (address, storage_diff) in self.storage_tries {
@@ -74,8 +80,16 @@ impl StorageTrieDiffEntry {
                     debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, ?database, "Difference in storage trie updates");
                 }
 
-                for (path, EntryDiff { task, regular, database }) in &storage_diff.removed_nodes {
-                    debug!(target: "engine::tree", ?address, ?path, ?task, ?regular, ?database, "Difference in removed account trie nodes");
+                for (
+                    path,
+                    EntryDiff {
+                        task: task_removed,
+                        regular: regular_removed,
+                        database: database_not_exists,
+                    },
+                ) in &storage_diff.removed_nodes
+                {
+                    debug!(target: "engine::tree", ?address, ?path, ?task_removed, ?regular_removed, ?database_not_exists, "Difference in removed storage trie nodes");
                 }
             }
         }
@@ -97,15 +111,13 @@ impl StorageTrieUpdatesDiff {
     }
 }
 
-/// Compares the trie updates from state root task and regular state root calculation, and logs
-/// the differences if there's any.
+/// Compares the trie updates from state root task, regular state root calculation and database,
+/// and logs the differences if there's any.
 pub(super) fn compare_trie_updates(
-    tx: &impl DbTx,
+    trie_cursor_factory: impl TrieCursorFactory,
     task: TrieUpdates,
     regular: TrieUpdates,
 ) -> Result<(), DatabaseError> {
-    let trie_cursor_factory = DatabaseTrieCursorFactory::new(tx);
-
     let mut task = adjust_trie_updates(task);
     let mut regular = adjust_trie_updates(regular);
 
@@ -118,7 +130,7 @@ pub(super) fn compare_trie_updates(
         .keys()
         .chain(regular.account_nodes.keys())
         .cloned()
-        .collect::<HashSet<_>>()
+        .collect::<BTreeSet<_>>()
     {
         let (task, regular) = (task.account_nodes.remove(&key), regular.account_nodes.remove(&key));
         let database = account_trie_cursor.seek_exact(key.clone())?.map(|x| x.1);
@@ -129,12 +141,13 @@ pub(super) fn compare_trie_updates(
     }
 
     // compare removed nodes
+    let mut account_trie_cursor = trie_cursor_factory.account_trie_cursor()?;
     for key in task
         .removed_nodes
         .iter()
         .chain(regular.removed_nodes.iter())
         .cloned()
-        .collect::<HashSet<_>>()
+        .collect::<BTreeSet<_>>()
     {
         let (task, regular) =
             (task.removed_nodes.contains(&key), regular.removed_nodes.contains(&key));
@@ -150,15 +163,17 @@ pub(super) fn compare_trie_updates(
         .keys()
         .chain(regular.storage_tries.keys())
         .copied()
-        .collect::<HashSet<_>>()
+        .collect::<BTreeSet<_>>()
     {
         let (mut task, mut regular) =
             (task.storage_tries.remove(&key), regular.storage_tries.remove(&key));
         if task != regular {
-            let mut storage_trie_cursor = trie_cursor_factory.storage_trie_cursor(key)?;
             if let Some((task, regular)) = task.as_mut().zip(regular.as_mut()) {
-                let storage_diff =
-                    compare_storage_trie_updates(&mut storage_trie_cursor, task, regular)?;
+                let storage_diff = compare_storage_trie_updates(
+                    || trie_cursor_factory.storage_trie_cursor(key),
+                    task,
+                    regular,
+                )?;
                 if storage_diff.has_differences() {
                     diff.storage_tries.insert(key, StorageTrieDiffEntry::Value(storage_diff));
                 }
@@ -177,12 +192,12 @@ pub(super) fn compare_trie_updates(
     Ok(())
 }
 
-fn compare_storage_trie_updates(
-    trie_cursor: &mut impl TrieCursor,
+fn compare_storage_trie_updates<C: TrieCursor>(
+    trie_cursor: impl Fn() -> Result<C, DatabaseError>,
     task: &mut StorageTrieUpdates,
     regular: &mut StorageTrieUpdates,
 ) -> Result<StorageTrieUpdatesDiff, DatabaseError> {
-    let database_deleted = trie_cursor.next()?.is_none();
+    let database_deleted = trie_cursor()?.next()?.is_none();
     let mut diff = StorageTrieUpdatesDiff {
         is_deleted: (task.is_deleted != regular.is_deleted).then_some(EntryDiff {
             task: task.is_deleted,
@@ -193,31 +208,33 @@ fn compare_storage_trie_updates(
     };
 
     // compare storage nodes
+    let mut storage_trie_cursor = trie_cursor()?;
     for key in task
         .storage_nodes
         .keys()
         .chain(regular.storage_nodes.keys())
         .cloned()
-        .collect::<HashSet<_>>()
+        .collect::<BTreeSet<_>>()
     {
         let (task, regular) = (task.storage_nodes.remove(&key), regular.storage_nodes.remove(&key));
-        let database = trie_cursor.seek_exact(key.clone())?.map(|x| x.1);
+        let database = storage_trie_cursor.seek_exact(key.clone())?.map(|x| x.1);
         if !branch_nodes_equal(task.as_ref(), regular.as_ref(), database.as_ref())? {
             diff.storage_nodes.insert(key, EntryDiff { task, regular, database });
         }
     }
 
     // compare removed nodes
+    let mut storage_trie_cursor = trie_cursor()?;
     for key in task
         .removed_nodes
         .iter()
         .chain(regular.removed_nodes.iter())
         .cloned()
-        .collect::<HashSet<_>>()
+        .collect::<BTreeSet<_>>()
     {
         let (task, regular) =
             (task.removed_nodes.contains(&key), regular.removed_nodes.contains(&key));
-        let database = trie_cursor.seek_exact(key.clone())?.map(|x| x.1).is_none();
+        let database = storage_trie_cursor.seek_exact(key.clone())?.map(|x| x.1).is_none();
         if task != regular {
             diff.removed_nodes.insert(key, EntryDiff { task, regular, database });
         }