From bacc61e567c1568bd19b6af52d54493978e716ee Mon Sep 17 00:00:00 2001
From: Federico Gimenez <fgimenez@users.noreply.github.com>
Date: Wed, 24 Jul 2024 16:29:39 +0200
Subject: [PATCH] feat: provide CanonicalInMemory state as input to
 EngineHandler (#9764)

---
 crates/engine/tree/src/tree/mod.rs                | 15 +++++++++------
 crates/ethereum/engine/src/service.rs             |  3 +++
 .../provider/src/providers/blockchain_provider.rs |  5 +++++
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/crates/engine/tree/src/tree/mod.rs b/crates/engine/tree/src/tree/mod.rs
index b953d3472..61cb5ed38 100644
--- a/crates/engine/tree/src/tree/mod.rs
+++ b/crates/engine/tree/src/tree/mod.rs
@@ -23,7 +23,7 @@ use reth_payload_primitives::{PayloadAttributes, PayloadBuilderAttributes, Paylo
 use reth_payload_validator::ExecutionPayloadValidator;
 use reth_primitives::{
     Block, BlockNumHash, BlockNumber, GotExpected, Header, Receipts, Requests, SealedBlock,
-    SealedBlockWithSenders, SealedHeader, B256, U256,
+    SealedBlockWithSenders, B256, U256,
 };
 use reth_provider::{
     BlockReader, ExecutionOutcome, StateProvider, StateProviderFactory, StateRootProvider,
@@ -353,7 +353,7 @@ where
         incoming: Receiver<FromEngine<BeaconEngineMessage<T>>>,
         outgoing: UnboundedSender<EngineApiEvent>,
         state: EngineApiTreeState,
-        header: SealedHeader,
+        canonical_in_memory_state: CanonicalInMemoryState,
         persistence: PersistenceHandle,
         payload_builder: PayloadBuilderHandle<T>,
     ) -> Self {
@@ -368,7 +368,7 @@ where
             persistence_state: PersistenceState::default(),
             is_backfill_active: false,
             state,
-            canonical_in_memory_state: CanonicalInMemoryState::with_head(header),
+            canonical_in_memory_state,
             payload_builder,
         }
     }
@@ -385,6 +385,7 @@ where
         incoming: Receiver<FromEngine<BeaconEngineMessage<T>>>,
         persistence: PersistenceHandle,
         payload_builder: PayloadBuilderHandle<T>,
+        canonical_in_memory_state: CanonicalInMemoryState,
     ) -> UnboundedReceiver<EngineApiEvent> {
         let best_block_number = provider.best_block_number().unwrap_or(0);
         let header = provider.sealed_header(best_block_number).ok().flatten().unwrap_or_default();
@@ -404,7 +405,7 @@ where
             incoming,
             tx,
             state,
-            header,
+            canonical_in_memory_state,
             persistence,
             payload_builder,
         );
@@ -1477,6 +1478,7 @@ mod tests {
 
             let header = chain_spec.genesis_header().seal_slow();
             let engine_api_tree_state = EngineApiTreeState::new(10, 10, header.num_hash());
+            let canonical_in_memory_state = CanonicalInMemoryState::with_head(header);
 
             let (to_payload_service, payload_command_rx) = unbounded_channel();
             let payload_builder = PayloadBuilderHandle::new(to_payload_service);
@@ -1488,7 +1490,7 @@ mod tests {
                 to_tree_rx,
                 from_tree_tx,
                 engine_api_tree_state,
-                header,
+                canonical_in_memory_state,
                 persistence_handle,
                 payload_builder,
             );
@@ -1544,6 +1546,7 @@ mod tests {
         };
 
         let header = blocks.first().unwrap().block().header.clone();
+        let canonical_in_memory_state = CanonicalInMemoryState::with_head(header);
 
         let (to_payload_service, payload_command_rx) = unbounded_channel();
         let payload_builder = PayloadBuilderHandle::new(to_payload_service);
@@ -1555,7 +1558,7 @@ mod tests {
             to_tree_rx,
             from_tree_tx,
             engine_api_tree_state,
-            header,
+            canonical_in_memory_state,
             persistence_handle,
             payload_builder,
         );
diff --git a/crates/ethereum/engine/src/service.rs b/crates/ethereum/engine/src/service.rs
index e9be329ba..37c36eee9 100644
--- a/crates/ethereum/engine/src/service.rs
+++ b/crates/ethereum/engine/src/service.rs
@@ -78,6 +78,8 @@ where
         let payload_validator = ExecutionPayloadValidator::new(chain_spec.clone());
         let executor_factory = EthExecutorProvider::ethereum(chain_spec);
 
+        let canonical_in_memory_state = blockchain_db.canonical_in_memory_state();
+
         let from_tree = EngineApiTreeHandlerImpl::spawn_new(
             blockchain_db,
             executor_factory,
@@ -86,6 +88,7 @@ where
             to_tree_rx,
             persistence_handle,
             payload_builder,
+            canonical_in_memory_state,
         );
 
         let engine_handler = EngineApiRequestHandler::new(to_tree_tx, from_tree);
diff --git a/crates/storage/provider/src/providers/blockchain_provider.rs b/crates/storage/provider/src/providers/blockchain_provider.rs
index f80ceb9c8..45cc5aa07 100644
--- a/crates/storage/provider/src/providers/blockchain_provider.rs
+++ b/crates/storage/provider/src/providers/blockchain_provider.rs
@@ -83,6 +83,11 @@ where
             None => Err(ProviderError::HeaderNotFound(best.best_number.into())),
         }
     }
+
+    /// Gets a clone of `canonical_in_memory_state`.
+    pub fn canonical_in_memory_state(&self) -> CanonicalInMemoryState {
+        self.canonical_in_memory_state.clone()
+    }
 }
 
 impl<DB> BlockchainProvider2<DB>