perf(engine): only clone headers instead of full blocks for tree tasks

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

@@ -1040,7 +1040,7 @@ where
             // we still need to process payload attributes if the head is already canonical
             if let Some(attr) = attrs {
                 let tip = self
-                    .block_header_by_hash(self.state.tree_state.canonical_block_hash())?
+                    .sealed_header_by_hash(self.state.tree_state.canonical_block_hash())?
                     .ok_or_else(|| {
                         // If we can't find the canonical block, then something is wrong and we need
                         // to return an error
@@ -1705,42 +1705,21 @@ where
         }))
     }
 
-    /// Return sealed block from database or in-memory state by hash.
+    /// Return sealed block header from in-memory state or database by hash.
     fn sealed_header_by_hash(
         &self,
         hash: B256,
     ) -> ProviderResult<Option<SealedHeader<N::BlockHeader>>> {
         // check memory first
-        let block = self
-            .state
-            .tree_state
-            .block_by_hash(hash)
-            .map(|block| block.as_ref().clone_sealed_header());
+        let header = self.state.tree_state.sealed_header_by_hash(&hash);
 
-        if block.is_some() {
-            Ok(block)
+        if header.is_some() {
+            Ok(header)
         } else {
             self.provider.sealed_header_by_hash(hash)
         }
     }
 
-    /// Return block header from database or in-memory state by hash.
-    fn block_header_by_hash(&self, hash: B256) -> ProviderResult<Option<N::BlockHeader>> {
-        // check database first
-        let mut header = self.provider.header_by_hash_or_number(hash.into())?;
-        if header.is_none() {
-            // Note: it's fine to return the unsealed block because the caller already has
-            // the hash
-            header = self
-                .state
-                .tree_state
-                .block_by_hash(hash)
-                // TODO: clone for compatibility. should we return an Arc here?
-                .map(|block| block.header().clone());
-        }
-        Ok(header)
-    }
-
     /// Return the parent hash of the lowest buffered ancestor for the requested block, if there
     /// are any buffered ancestors. If there are no buffered ancestors, and the block itself does
     /// not exist in the buffer, this returns the hash that is passed in.
@@ -1770,7 +1749,7 @@ where
         parent_hash: B256,
     ) -> ProviderResult<Option<B256>> {
         // Check if parent exists in side chain or in canonical chain.
-        if self.block_header_by_hash(parent_hash)?.is_some() {
+        if self.sealed_header_by_hash(parent_hash)?.is_some() {
             return Ok(Some(parent_hash))
         }
 
@@ -1784,7 +1763,7 @@ where
 
             // If current_header is None, then the current_hash does not have an invalid
             // ancestor in the cache, check its presence in blockchain tree
-            if current_block.is_none() && self.block_header_by_hash(current_hash)?.is_some() {
+            if current_block.is_none() && self.sealed_header_by_hash(current_hash)?.is_some() {
                 return Ok(Some(current_hash))
             }
         }
@@ -1797,7 +1776,7 @@ where
     fn prepare_invalid_response(&mut self, mut parent_hash: B256) -> ProviderResult<PayloadStatus> {
         // Edge case: the `latestValid` field is the zero hash if the parent block is the terminal
         // PoW block, which we need to identify by looking at the parent's block difficulty
-        if let Some(parent) = self.block_header_by_hash(parent_hash)? {
+        if let Some(parent) = self.sealed_header_by_hash(parent_hash)? {
             if !parent.difficulty().is_zero() {
                 parent_hash = B256::ZERO;
             }
@@ -2301,7 +2280,7 @@ where
         let block_num_hash = block_id.block;
         debug!(target: "engine::tree", block=?block_num_hash, parent = ?block_id.parent, "Inserting new block into tree");
 
-        match self.block_header_by_hash(block_num_hash.hash) {
+        match self.sealed_header_by_hash(block_num_hash.hash) {
             Err(err) => {
                 let block = convert_to_block(self, input)?;
                 return Err(InsertBlockError::new(block.into_sealed_block(), err.into()).into());

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

@@ -606,18 +606,17 @@ where
         })
     }
 
-    /// Return sealed block from database or in-memory state by hash.
+    /// Return sealed block header from database or in-memory state by hash.
     fn sealed_header_by_hash(
         &self,
         hash: B256,
         state: &EngineApiTreeState<N>,
     ) -> ProviderResult<Option<SealedHeader<N::BlockHeader>>> {
         // check memory first
-        let block =
-            state.tree_state.block_by_hash(hash).map(|block| block.as_ref().clone_sealed_header());
+        let header = state.tree_state.sealed_header_by_hash(&hash);
 
-        if block.is_some() {
-            Ok(block)
+        if header.is_some() {
+            Ok(header)
         } else {
             self.provider.sealed_header_by_hash(hash)
         }

crates/engine/tree/src/tree/state.rs

@@ -7,7 +7,7 @@ use alloy_primitives::{
     BlockNumber, B256,
 };
 use reth_chain_state::{EthPrimitives, ExecutedBlockWithTrieUpdates};
-use reth_primitives_traits::{AlloyBlockHeader, NodePrimitives, SealedBlock};
+use reth_primitives_traits::{AlloyBlockHeader, NodePrimitives, SealedHeader};
 use reth_trie::updates::TrieUpdates;
 use std::{
     collections::{btree_map, hash_map, BTreeMap, VecDeque},
@@ -85,9 +85,12 @@ impl<N: NodePrimitives> TreeState<N> {
         self.blocks_by_hash.get(&hash)
     }
 
-    /// Returns the block by hash.
-    pub(crate) fn block_by_hash(&self, hash: B256) -> Option<Arc<SealedBlock<N::Block>>> {
-        self.blocks_by_hash.get(&hash).map(|b| Arc::new(b.recovered_block().sealed_block().clone()))
+    /// Returns the sealed block header by hash.
+    pub(crate) fn sealed_header_by_hash(
+        &self,
+        hash: &B256,
+    ) -> Option<SealedHeader<N::BlockHeader>> {
+        self.blocks_by_hash.get(hash).map(|b| b.sealed_block().sealed_header().clone())
     }
 
     /// Returns all available blocks for the given hash that lead back to the canonical chain, from

crates/engine/tree/src/tree/tests.rs

@@ -760,7 +760,7 @@ async fn test_get_canonical_blocks_to_persist() {
     let fork_block_hash = fork_block.recovered_block().hash();
     test_harness.tree.state.tree_state.insert_executed(fork_block);
 
-    assert!(test_harness.tree.state.tree_state.block_by_hash(fork_block_hash).is_some());
+    assert!(test_harness.tree.state.tree_state.sealed_header_by_hash(&fork_block_hash).is_some());
 
     let blocks_to_persist = test_harness.tree.get_canonical_blocks_to_persist().unwrap();
     assert_eq!(blocks_to_persist.len(), expected_blocks_to_persist_length);