fix: gracefully handle missing persisted_trie_updates (#13942)

crates/chain-state/src/in_memory.rs

@@ -240,7 +240,7 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
     /// Updates the pending block with the given block.
     ///
     /// Note: This assumes that the parent block of the pending block is canonical.
-    pub fn set_pending_block(&self, pending: ExecutedBlock<N>) {
+    pub fn set_pending_block(&self, pending: ExecutedBlockWithTrieUpdates<N>) {
         // fetch the state of the pending block's parent block
         let parent = self.state_by_hash(pending.recovered_block().parent_hash());
         let pending = BlockState::with_parent(pending, parent);
@@ -254,9 +254,10 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
     ///
     /// This removes all reorged blocks and appends the new blocks to the tracked chain and connects
     /// them to their parent blocks.
-    fn update_blocks<I>(&self, new_blocks: I, reorged: I)
+    fn update_blocks<I, R>(&self, new_blocks: I, reorged: R)
     where
-        I: IntoIterator<Item = ExecutedBlock<N>>,
+        I: IntoIterator<Item = ExecutedBlockWithTrieUpdates<N>>,
+        R: IntoIterator<Item = ExecutedBlock<N>>,
     {
         {
             // acquire locks, starting with the numbers lock
@@ -601,7 +602,7 @@ impl<N: NodePrimitives> CanonicalInMemoryState<N> {
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct BlockState<N: NodePrimitives = EthPrimitives> {
     /// The executed block that determines the state after this block has been executed.
-    block: ExecutedBlock<N>,
+    block: ExecutedBlockWithTrieUpdates<N>,
     /// The block's parent block if it exists.
     parent: Option<Arc<BlockState<N>>>,
 }
@@ -609,12 +610,15 @@ pub struct BlockState<N: NodePrimitives = EthPrimitives> {
 #[allow(dead_code)]
 impl<N: NodePrimitives> BlockState<N> {
     /// [`BlockState`] constructor.
-    pub const fn new(block: ExecutedBlock<N>) -> Self {
+    pub const fn new(block: ExecutedBlockWithTrieUpdates<N>) -> Self {
         Self { block, parent: None }
     }
 
     /// [`BlockState`] constructor with parent.
-    pub const fn with_parent(block: ExecutedBlock<N>, parent: Option<Arc<Self>>) -> Self {
+    pub const fn with_parent(
+        block: ExecutedBlockWithTrieUpdates<N>,
+        parent: Option<Arc<Self>>,
+    ) -> Self {
         Self { block, parent }
     }
 
@@ -628,12 +632,12 @@ impl<N: NodePrimitives> BlockState<N> {
     }
 
     /// Returns the executed block that determines the state.
-    pub fn block(&self) -> ExecutedBlock<N> {
+    pub fn block(&self) -> ExecutedBlockWithTrieUpdates<N> {
         self.block.clone()
     }
 
     /// Returns a reference to the executed block that determines the state.
-    pub const fn block_ref(&self) -> &ExecutedBlock<N> {
+    pub const fn block_ref(&self) -> &ExecutedBlockWithTrieUpdates<N> {
         &self.block
     }
 
@@ -787,7 +791,7 @@ impl<N: NodePrimitives> BlockState<N> {
 }
 
 /// Represents an executed block stored in-memory.
-#[derive(Clone, Debug, PartialEq, Eq, Default)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub struct ExecutedBlock<N: NodePrimitives = EthPrimitives> {
     /// Recovered Block
     pub recovered_block: Arc<RecoveredBlock<N::Block>>,
@@ -795,21 +799,19 @@ pub struct ExecutedBlock<N: NodePrimitives = EthPrimitives> {
     pub execution_output: Arc<ExecutionOutcome<N::Receipt>>,
     /// Block's hashed state.
     pub hashed_state: Arc<HashedPostState>,
-    /// Trie updates that result of applying the block.
-    pub trie: Arc<TrieUpdates>,
+}
+
+impl<N: NodePrimitives> Default for ExecutedBlock<N> {
+    fn default() -> Self {
+        Self {
+            recovered_block: Default::default(),
+            execution_output: Default::default(),
+            hashed_state: Default::default(),
+        }
+    }
 }
 
 impl<N: NodePrimitives> ExecutedBlock<N> {
-    /// [`ExecutedBlock`] constructor.
-    pub const fn new(
-        recovered_block: Arc<RecoveredBlock<N::Block>>,
-        execution_output: Arc<ExecutionOutcome<N::Receipt>>,
-        hashed_state: Arc<HashedPostState>,
-        trie: Arc<TrieUpdates>,
-    ) -> Self {
-        Self { recovered_block, execution_output, hashed_state, trie }
-    }
-
     /// Returns a reference to an inner [`SealedBlock`]
     #[inline]
     pub fn sealed_block(&self) -> &SealedBlock<N::Block> {
@@ -833,6 +835,42 @@ impl<N: NodePrimitives> ExecutedBlock<N> {
     pub fn hashed_state(&self) -> &HashedPostState {
         &self.hashed_state
     }
+}
+
+/// An [`ExecutedBlock`] with its [`TrieUpdates`].
+///
+/// We store it as separate type because [`TrieUpdates`] are only available for blocks stored in
+/// memory and can't be obtained for canonical persisted blocks.
+#[derive(
+    Clone,
+    Debug,
+    PartialEq,
+    Eq,
+    Default,
+    derive_more::Deref,
+    derive_more::DerefMut,
+    derive_more::Into,
+)]
+pub struct ExecutedBlockWithTrieUpdates<N: NodePrimitives = EthPrimitives> {
+    /// Inner [`ExecutedBlock`].
+    #[deref]
+    #[deref_mut]
+    #[into]
+    pub block: ExecutedBlock<N>,
+    /// Trie updates that result of applying the block.
+    pub trie: Arc<TrieUpdates>,
+}
+
+impl<N: NodePrimitives> ExecutedBlockWithTrieUpdates<N> {
+    /// [`ExecutedBlock`] constructor.
+    pub const fn new(
+        recovered_block: Arc<RecoveredBlock<N::Block>>,
+        execution_output: Arc<ExecutionOutcome<N::Receipt>>,
+        hashed_state: Arc<HashedPostState>,
+        trie: Arc<TrieUpdates>,
+    ) -> Self {
+        Self { block: ExecutedBlock { recovered_block, execution_output, hashed_state }, trie }
+    }
 
     /// Returns a reference to the trie updates for the block
     #[inline]
@@ -847,14 +885,18 @@ pub enum NewCanonicalChain<N: NodePrimitives = EthPrimitives> {
     /// A simple append to the current canonical head
     Commit {
         /// all blocks that lead back to the canonical head
-        new: Vec<ExecutedBlock<N>>,
+        new: Vec<ExecutedBlockWithTrieUpdates<N>>,
     },
     /// A reorged chain consists of two chains that trace back to a shared ancestor block at which
     /// point they diverge.
     Reorg {
         /// All blocks of the _new_ chain
-        new: Vec<ExecutedBlock<N>>,
+        new: Vec<ExecutedBlockWithTrieUpdates<N>>,
         /// All blocks of the _old_ chain
+        ///
+        /// These are not [`ExecutedBlockWithTrieUpdates`] because we don't always have the trie
+        /// updates for the old canonical chain. For example, in case of node being restarted right
+        /// before the reorg [`TrieUpdates`] can't be fetched from database.
         old: Vec<ExecutedBlock<N>>,
     },
 }
@@ -1257,7 +1299,7 @@ mod tests {
             block1.recovered_block().hash()
         );
 
-        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1] };
+        let chain = NewCanonicalChain::Reorg { new: vec![block2.clone()], old: vec![block1.block] };
         state.update_chain(chain);
         assert_eq!(
             state.head_state().unwrap().block_ref().recovered_block().hash(),
@@ -1540,7 +1582,7 @@ mod tests {
         // Test reorg notification
         let chain_reorg = NewCanonicalChain::Reorg {
             new: vec![block1a.clone(), block2a.clone()],
-            old: vec![block1.clone(), block2.clone()],
+            old: vec![block1.block.clone(), block2.block.clone()],
         };
 
         assert_eq!(

crates/chain-state/src/memory_overlay.rs

@@ -1,4 +1,4 @@
-use super::ExecutedBlock;
+use super::ExecutedBlockWithTrieUpdates;
 use alloy_consensus::BlockHeader;
 use alloy_primitives::{
     keccak256, map::B256HashMap, Address, BlockNumber, Bytes, StorageKey, StorageValue, B256,
@@ -23,7 +23,7 @@ pub struct MemoryOverlayStateProviderRef<'a, N: NodePrimitives = reth_primitives
     /// Historical state provider for state lookups that are not found in in-memory blocks.
     pub(crate) historical: Box<dyn StateProvider + 'a>,
     /// The collection of executed parent blocks. Expected order is newest to oldest.
-    pub(crate) in_memory: Vec<ExecutedBlock<N>>,
+    pub(crate) in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
     /// Lazy-loaded in-memory trie data.
     pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
 }
@@ -40,7 +40,10 @@ impl<'a, N: NodePrimitives> MemoryOverlayStateProviderRef<'a, N> {
     /// - `in_memory` - the collection of executed ancestor blocks in reverse.
     /// - `historical` - a historical state provider for the latest ancestor block stored in the
     ///   database.
-    pub fn new(historical: Box<dyn StateProvider + 'a>, in_memory: Vec<ExecutedBlock<N>>) -> Self {
+    pub fn new(
+        historical: Box<dyn StateProvider + 'a>,
+        in_memory: Vec<ExecutedBlockWithTrieUpdates<N>>,
+    ) -> Self {
         Self { historical, in_memory, trie_state: OnceLock::new() }
     }
 

crates/chain-state/src/test_utils.rs

@@ -1,7 +1,7 @@
 use core::marker::PhantomData;
 
 use crate::{
-    in_memory::ExecutedBlock, CanonStateNotification, CanonStateNotifications,
+    in_memory::ExecutedBlockWithTrieUpdates, CanonStateNotification, CanonStateNotifications,
     CanonStateSubscriptions,
 };
 use alloy_consensus::{
@@ -204,17 +204,17 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
         fork
     }
 
-    /// Gets an [`ExecutedBlock`] with [`BlockNumber`], [`Receipts`] and parent hash.
+    /// Gets an [`ExecutedBlockWithTrieUpdates`] with [`BlockNumber`], [`Receipts`] and parent hash.
     fn get_executed_block(
         &mut self,
         block_number: BlockNumber,
         receipts: Receipts,
         parent_hash: B256,
-    ) -> ExecutedBlock {
+    ) -> ExecutedBlockWithTrieUpdates {
         let block_with_senders = self.generate_random_block(block_number, parent_hash);
 
         let (block, senders) = block_with_senders.split_sealed();
-        ExecutedBlock::new(
+        ExecutedBlockWithTrieUpdates::new(
             Arc::new(RecoveredBlock::new_sealed(block, senders)),
             Arc::new(ExecutionOutcome::new(
                 BundleState::default(),
@@ -227,22 +227,22 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
         )
     }
 
-    /// Generates an [`ExecutedBlock`] that includes the given [`Receipts`].
+    /// Generates an [`ExecutedBlockWithTrieUpdates`] that includes the given [`Receipts`].
     pub fn get_executed_block_with_receipts(
         &mut self,
         receipts: Receipts,
         parent_hash: B256,
-    ) -> ExecutedBlock {
+    ) -> ExecutedBlockWithTrieUpdates {
         let number = rand::thread_rng().gen::<u64>();
         self.get_executed_block(number, receipts, parent_hash)
     }
 
-    /// Generates an [`ExecutedBlock`] with the given [`BlockNumber`].
+    /// Generates an [`ExecutedBlockWithTrieUpdates`] with the given [`BlockNumber`].
     pub fn get_executed_block_with_number(
         &mut self,
         block_number: BlockNumber,
         parent_hash: B256,
-    ) -> ExecutedBlock {
+    ) -> ExecutedBlockWithTrieUpdates {
         self.get_executed_block(block_number, Receipts { receipt_vec: vec![vec![]] }, parent_hash)
     }
 
@@ -250,7 +250,7 @@ impl<N: NodePrimitives> TestBlockBuilder<N> {
     pub fn get_executed_blocks(
         &mut self,
         range: Range<u64>,
-    ) -> impl Iterator<Item = ExecutedBlock> + '_ {
+    ) -> impl Iterator<Item = ExecutedBlockWithTrieUpdates> + '_ {
         let mut parent_hash = B256::default();
         range.map(move |number| {
             let current_parent_hash = parent_hash;