refactor(engine): persistence logic

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

@@ -1568,6 +1568,9 @@ where
     /// `(last_persisted_number .. canonical_head - threshold]`. The expected
     /// order is oldest -> newest.
     ///
+    /// If any blocks are missing trie updates, all blocks are persisted, not taking `threshold`
+    /// into account.
+    ///
     /// For those blocks that didn't have the trie updates calculated, runs the state root
     /// calculation, and saves the trie updates.
     ///
@@ -1582,13 +1585,31 @@ where
         let mut blocks_to_persist = Vec::new();
         let mut current_hash = self.state.tree_state.canonical_block_hash();
         let last_persisted_number = self.persistence_state.last_persisted_block.number;
-
         let canonical_head_number = self.state.tree_state.canonical_block_number();
+        let all_blocks_have_trie_updates = self
+            .state
+            .tree_state
+            .blocks_by_hash
+            .values()
+            .all(|block| block.trie_updates().is_some());
 
-        let target_number =
-            canonical_head_number.saturating_sub(self.config.memory_block_buffer_target());
+        let target_number = if all_blocks_have_trie_updates {
+            // Persist only up to block buffer target if all blocks have trie updates
+            canonical_head_number.saturating_sub(self.config.memory_block_buffer_target())
+        } else {
+            // Persist all blocks if any block is missing trie updates
+            canonical_head_number
+        };
 
-        debug!(target: "engine::tree", ?last_persisted_number, ?canonical_head_number, ?target_number, ?current_hash, "Returning canonical blocks to persist");
+        debug!(
+            target: "engine::tree",
+            ?current_hash,
+            ?last_persisted_number,
+            ?canonical_head_number,
+            ?all_blocks_have_trie_updates,
+            ?target_number,
+            "Returning canonical blocks to persist"
+        );
         while let Some(block) = self.state.tree_state.blocks_by_hash.get(&current_hash) {
             if block.recovered_block().number() <= last_persisted_number {
                 break;
@@ -2332,12 +2353,8 @@ where
             Ok(is_fork) => is_fork,
         };
 
-        let ctx = TreeCtx::new(
-            &mut self.state,
-            &self.persistence_state,
-            &self.canonical_in_memory_state,
-            is_fork,
-        );
+        let ctx =
+            TreeCtx::new(&mut self.state, &self.persistence_state, &self.canonical_in_memory_state);
 
         let start = Instant::now();
 

crates/engine/tree/src/tree/payload_validator.rs

@@ -35,9 +35,9 @@ use reth_primitives_traits::{
     AlloyBlockHeader, BlockTy, GotExpected, NodePrimitives, RecoveredBlock, SealedHeader,
 };
 use reth_provider::{
-    BlockExecutionOutput, BlockNumReader, BlockReader, DBProvider, DatabaseProviderFactory,
-    ExecutionOutcome, HashedPostStateProvider, ProviderError, StateProvider, StateProviderFactory,
-    StateReader, StateRootProvider,
+    BlockExecutionOutput, BlockHashReader, BlockNumReader, BlockReader, DBProvider,
+    DatabaseProviderFactory, ExecutionOutcome, HashedPostStateProvider, HeaderProvider,
+    ProviderError, StateProvider, StateProviderFactory, StateReader, StateRootProvider,
 };
 use reth_revm::db::State;
 use reth_trie::{updates::TrieUpdates, HashedPostState, KeccakKeyHasher, TrieInput};
@@ -57,8 +57,6 @@ pub struct TreeCtx<'a, N: NodePrimitives> {
     persistence: &'a PersistenceState,
     /// Reference to the canonical in-memory state
     canonical_in_memory_state: &'a CanonicalInMemoryState<N>,
-    /// Whether the currently validated block is on a fork chain.
-    is_fork: bool,
 }
 
 impl<'a, N: NodePrimitives> std::fmt::Debug for TreeCtx<'a, N> {
@@ -77,9 +75,8 @@ impl<'a, N: NodePrimitives> TreeCtx<'a, N> {
         state: &'a mut EngineApiTreeState<N>,
         persistence: &'a PersistenceState,
         canonical_in_memory_state: &'a CanonicalInMemoryState<N>,
-        is_fork: bool,
     ) -> Self {
-        Self { state, persistence, canonical_in_memory_state, is_fork }
+        Self { state, persistence, canonical_in_memory_state }
     }
 
     /// Returns a reference to the engine tree state
@@ -102,11 +99,6 @@ impl<'a, N: NodePrimitives> TreeCtx<'a, N> {
         self.canonical_in_memory_state
     }
 
-    /// Returns whether the currently validated block is on a fork chain.
-    pub const fn is_fork(&self) -> bool {
-        self.is_fork
-    }
-
     /// Determines the persisting kind for the given block based on persistence info.
     ///
     /// Based on the given header it returns whether any conflicting persistence operation is
@@ -588,9 +580,26 @@ where
         // terminate prewarming task with good state output
         handle.terminate_caching(Some(output.state.clone()));
 
-        // If the block is a fork, we don't save the trie updates, because they may be incorrect.
+        // If the block doesn't connect to the database tip, we don't save its trie updates, because
+        // they may be incorrect as they were calculated on top of the forked block.
+        //
+        // We also only save trie updates if all ancestors have trie updates, because otherwise the
+        // trie updates may be incorrect.
+        //
         // Instead, they will be recomputed on persistence.
-        let trie_updates = if ctx.is_fork() {
+        let connects_to_last_persisted =
+            ensure_ok!(self.block_connects_to_last_persisted(ctx, &block));
+        let should_discard_trie_updates =
+            !connects_to_last_persisted || has_ancestors_with_missing_trie_updates;
+        debug!(
+            target: "engine::tree",
+            block = ?block_num_hash,
+            connects_to_last_persisted,
+            has_ancestors_with_missing_trie_updates,
+            should_discard_trie_updates,
+            "Checking if should discard trie updates"
+        );
+        let trie_updates = if should_discard_trie_updates {
             ExecutedTrieUpdates::Missing
         } else {
             ExecutedTrieUpdates::Present(Arc::new(trie_output))
@@ -729,6 +738,51 @@ where
         ParallelStateRoot::new(consistent_view, input).incremental_root_with_updates()
     }
 
+    /// Checks if the given block connects to the last persisted block, i.e. if the last persisted
+    /// block is the ancestor of the given block.
+    ///
+    /// This checks the database for the actual last persisted block, not [`PersistenceState`].
+    fn block_connects_to_last_persisted(
+        &self,
+        ctx: TreeCtx<'_, N>,
+        block: &RecoveredBlock<N::Block>,
+    ) -> ProviderResult<bool> {
+        let provider = self.provider.database_provider_ro()?;
+        let last_persisted_block = provider.last_block_number()?;
+        let last_persisted_hash = provider
+            .block_hash(last_persisted_block)?
+            .ok_or(ProviderError::HeaderNotFound(last_persisted_block.into()))?;
+        let last_persisted = NumHash::new(last_persisted_block, last_persisted_hash);
+
+        let parent_num_hash = |hash: B256| -> ProviderResult<NumHash> {
+            let parent_num_hash =
+                if let Some(header) = ctx.state().tree_state.sealed_header_by_hash(&hash) {
+                    Some(header.parent_num_hash())
+                } else {
+                    provider.sealed_header_by_hash(hash)?.map(|header| header.parent_num_hash())
+                };
+
+            parent_num_hash.ok_or(ProviderError::BlockHashNotFound(hash))
+        };
+
+        let mut parent_block = block.parent_num_hash();
+        while parent_block.number > last_persisted.number {
+            parent_block = parent_num_hash(parent_block.hash)?;
+        }
+
+        let connects = parent_block == last_persisted;
+
+        debug!(
+            target: "engine::tree",
+            num_hash = ?block.num_hash(),
+            ?last_persisted,
+            ?parent_block,
+            "Checking if block connects to last persisted block"
+        );
+
+        Ok(connects)
+    }
+
     /// Check if the given block has any ancestors with missing trie updates.
     fn has_ancestors_with_missing_trie_updates(
         &self,