perf(engine): lazy load overlay trie state (#10541)

Author: rkrasiuk

crates/chain-state/src/in_memory.rs

@@ -493,7 +493,7 @@ impl CanonicalInMemoryState {
             Vec::new()
         };
 
-        MemoryOverlayStateProvider::new(in_memory, historical)
+        MemoryOverlayStateProvider::new(historical, in_memory)
     }
 
     /// Returns an iterator over all canonical blocks in the in-memory state, from newest to oldest.

crates/chain-state/src/memory_overlay.rs

@@ -11,20 +11,18 @@ use reth_trie::{
     prefix_set::TriePrefixSetsMut, updates::TrieUpdates, AccountProof, HashedPostState,
     HashedStorage,
 };
-use std::collections::HashMap;
+use std::{collections::HashMap, sync::OnceLock};
 
 /// A state provider that stores references to in-memory blocks along with their state as well as
 /// the historical state provider for fallback lookups.
 #[allow(missing_debug_implementations)]
 pub struct MemoryOverlayStateProvider {
-    /// The collection of executed parent blocks. Expected order is newest to oldest.
-    pub(crate) in_memory: Vec<ExecutedBlock>,
-    /// The collection of hashed state from in-memory blocks.
-    pub(crate) hashed_post_state: HashedPostState,
-    /// The collection of aggregated in-memory trie updates.
-    pub(crate) trie_updates: TrieUpdates,
     /// Historical state provider for state lookups that are not found in in-memory blocks.
     pub(crate) historical: Box<dyn StateProvider>,
+    /// The collection of executed parent blocks. Expected order is newest to oldest.
+    pub(crate) in_memory: Vec<ExecutedBlock>,
+    /// Lazy-loaded in-memory trie data.
+    pub(crate) trie_state: OnceLock<MemoryOverlayTrieState>,
 }
 
 impl MemoryOverlayStateProvider {
@@ -35,20 +33,29 @@ impl MemoryOverlayStateProvider {
     /// - `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(in_memory: Vec<ExecutedBlock>, historical: Box<dyn StateProvider>) -> Self {
-        let mut hashed_post_state = HashedPostState::default();
-        let mut trie_updates = TrieUpdates::default();
-        for block in in_memory.iter().rev() {
-            hashed_post_state.extend_ref(block.hashed_state.as_ref());
-            trie_updates.extend_ref(block.trie.as_ref());
-        }
-        Self { in_memory, hashed_post_state, trie_updates, historical }
+    pub fn new(historical: Box<dyn StateProvider>, in_memory: Vec<ExecutedBlock>) -> Self {
+        Self { historical, in_memory, trie_state: OnceLock::new() }
     }
 
     /// Turn this state provider into a [`StateProviderBox`]
     pub fn boxed(self) -> StateProviderBox {
         Box::new(self)
     }
+
+    /// Return lazy-loaded trie state aggregated from in-memory blocks.
+    fn trie_state(&self) -> MemoryOverlayTrieState {
+        self.trie_state
+            .get_or_init(|| {
+                let mut hashed_state = HashedPostState::default();
+                let mut trie_nodes = TrieUpdates::default();
+                for block in self.in_memory.iter().rev() {
+                    hashed_state.extend_ref(block.hashed_state.as_ref());
+                    trie_nodes.extend_ref(block.trie.as_ref());
+                }
+                MemoryOverlayTrieState { trie_nodes, hashed_state }
+            })
+            .clone()
+    }
 }
 
 impl BlockHashReader for MemoryOverlayStateProvider {
@@ -105,14 +112,13 @@ impl StateRootProvider for MemoryOverlayStateProvider {
     fn hashed_state_root_from_nodes(
         &self,
         nodes: TrieUpdates,
-        hashed_state: HashedPostState,
+        state: HashedPostState,
         prefix_sets: TriePrefixSetsMut,
     ) -> ProviderResult<B256> {
-        let mut trie_nodes = self.trie_updates.clone();
+        let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state();
         trie_nodes.extend(nodes);
-        let mut state = self.hashed_post_state.clone();
-        state.extend(hashed_state);
-        self.historical.hashed_state_root_from_nodes(trie_nodes, state, prefix_sets)
+        hashed_state.extend(state);
+        self.historical.hashed_state_root_from_nodes(trie_nodes, hashed_state, prefix_sets)
     }
 
     fn hashed_state_root_with_updates(
@@ -130,14 +136,17 @@ impl StateRootProvider for MemoryOverlayStateProvider {
     fn hashed_state_root_from_nodes_with_updates(
         &self,
         nodes: TrieUpdates,
-        hashed_state: HashedPostState,
+        state: HashedPostState,
         prefix_sets: TriePrefixSetsMut,
     ) -> ProviderResult<(B256, TrieUpdates)> {
-        let mut trie_nodes = self.trie_updates.clone();
+        let MemoryOverlayTrieState { mut trie_nodes, mut hashed_state } = self.trie_state();
         trie_nodes.extend(nodes);
-        let mut state = self.hashed_post_state.clone();
-        state.extend(hashed_state);
-        self.historical.hashed_state_root_from_nodes_with_updates(trie_nodes, state, prefix_sets)
+        hashed_state.extend(state);
+        self.historical.hashed_state_root_from_nodes_with_updates(
+            trie_nodes,
+            hashed_state,
+            prefix_sets,
+        )
     }
 
     // TODO: Currently this does not reuse available in-memory trie nodes.
@@ -153,13 +162,13 @@ impl StateRootProvider for MemoryOverlayStateProvider {
 impl StateProofProvider for MemoryOverlayStateProvider {
     fn hashed_proof(
         &self,
-        hashed_state: HashedPostState,
+        state: HashedPostState,
         address: Address,
         slots: &[B256],
     ) -> ProviderResult<AccountProof> {
-        let mut state = self.hashed_post_state.clone();
-        state.extend(hashed_state);
-        self.historical.hashed_proof(state, address, slots)
+        let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state();
+        hashed_state.extend(state);
+        self.historical.hashed_proof(hashed_state, address, slots)
     }
 
     // TODO: Currently this does not reuse available in-memory trie nodes.
@@ -168,9 +177,9 @@ impl StateProofProvider for MemoryOverlayStateProvider {
         overlay: HashedPostState,
         target: HashedPostState,
     ) -> ProviderResult<HashMap<B256, Bytes>> {
-        let mut state = self.hashed_post_state.clone();
-        state.extend(overlay);
-        self.historical.witness(state, target)
+        let MemoryOverlayTrieState { mut hashed_state, .. } = self.trie_state();
+        hashed_state.extend(overlay);
+        self.historical.witness(hashed_state, target)
     }
 }
 
@@ -199,3 +208,12 @@ impl StateProvider for MemoryOverlayStateProvider {
         self.historical.bytecode_by_hash(code_hash)
     }
 }
+
+/// The collection of data necessary for trie-related operations for [`MemoryOverlayStateProvider`].
+#[derive(Clone, Debug)]
+pub(crate) struct MemoryOverlayTrieState {
+    /// The collection of aggregated in-memory trie updates.
+    pub(crate) trie_nodes: TrieUpdates,
+    /// The collection of hashed state from in-memory blocks.
+    pub(crate) hashed_state: HashedPostState,
+}

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

@@ -1232,7 +1232,7 @@ where
             trace!(target: "engine", %hash, "found canonical state for block in memory");
             // the block leads back to the canonical chain
             let historical = self.provider.state_by_block_hash(historical)?;
-            return Ok(Some(Box::new(MemoryOverlayStateProvider::new(blocks, historical))))
+            return Ok(Some(Box::new(MemoryOverlayStateProvider::new(historical, blocks))))
         }
 
         // the hash could belong to an unknown block or a persisted block