feat: async header response transform (#335)

Author: frisitano

Cargo.lock

@@ -60,6 +60,7 @@ reth-metrics = { workspace = true, features = ["common"] }
 metrics.workspace = true
 
 # misc
+async-trait.workspace = true
 auto_impl.workspace = true
 aquamarine.workspace = true
 tracing.workspace = true

crates/net/network/Cargo.toml

@@ -9,12 +9,13 @@ use crate::{
         policy::NetworkPolicies,
         TransactionPropagationPolicy, TransactionsManager, TransactionsManagerConfig,
     },
-    transform::header::HeaderTransform,
+    transform::header::HeaderResponseTransform,
     NetworkHandle, NetworkManager,
 };
 use reth_eth_wire::{EthNetworkPrimitives, NetworkPrimitives};
 use reth_network_api::test_utils::PeersHandleProvider;
 use reth_transaction_pool::TransactionPool;
+use std::sync::Arc;
 use tokio::sync::mpsc;
 
 /// We set the max channel capacity of the `EthRequestHandler` to 256
@@ -64,7 +65,7 @@ impl<Tx, Eth, N: NetworkPrimitives> NetworkBuilder<Tx, Eth, N> {
     pub fn request_handler<Client>(
         self,
         client: Client,
-        request_header_transform: Option<Box<dyn HeaderTransform<N::BlockHeader>>>,
+        request_header_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
     ) -> NetworkBuilder<Tx, EthRequestHandler<Client, N>, N> {
         let Self { mut network, transactions, .. } = self;
         let (tx, rx) = mpsc::channel(ETH_REQUEST_CHANNEL_CAPACITY);

crates/net/network/src/builder.rs

@@ -2,12 +2,12 @@
 
 use crate::{
     budget::DEFAULT_BUDGET_TRY_DRAIN_DOWNLOADERS, metered_poll_nested_stream_with_budget,
-    metrics::EthRequestHandlerMetrics, transform::header::HeaderTransform,
+    metrics::EthRequestHandlerMetrics, transform::header::HeaderResponseTransform,
 };
 use alloy_consensus::{BlockHeader, ReceiptWithBloom};
 use alloy_eips::BlockHashOrNumber;
 use alloy_rlp::Encodable;
-use futures::StreamExt;
+use futures::{future::join_all, StreamExt};
 use reth_eth_wire::{
     BlockBodies, BlockHeaders, EthNetworkPrimitives, GetBlockBodies, GetBlockHeaders, GetNodeData,
     GetReceipts, HeadersDirection, NetworkPrimitives, NodeData, Receipts, Receipts69,
@@ -20,6 +20,7 @@ use reth_storage_api::{BlockReader, HeaderProvider};
 use std::{
     future::Future,
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
     time::Duration,
 };
@@ -54,15 +55,15 @@ pub const SOFT_RESPONSE_LIMIT: usize = 2 * 1024 * 1024;
 #[must_use = "Manager does nothing unless polled."]
 pub struct EthRequestHandler<C, N: NetworkPrimitives = EthNetworkPrimitives> {
     /// The client type that can interact with the chain.
-    client: C,
+    client: Arc<C>,
     /// Used for reporting peers.
     // TODO use to report spammers
     #[expect(dead_code)]
     peers: PeersHandle,
     /// Incoming request from the [`NetworkManager`](crate::NetworkManager).
     incoming_requests: ReceiverStream<IncomingEthRequest<N>>,
     /// The header transform to apply to the headers before sending to peers.
-    header_transform: Option<Box<dyn HeaderTransform<N::BlockHeader>>>,
+    header_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
     /// Metrics for the eth request handler.
     metrics: EthRequestHandlerMetrics,
 }
@@ -74,10 +75,10 @@ impl<C, N: NetworkPrimitives> EthRequestHandler<C, N> {
         client: C,
         peers: PeersHandle,
         incoming: Receiver<IncomingEthRequest<N>>,
-        header_transform: Option<Box<dyn HeaderTransform<N::BlockHeader>>>,
+        header_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
     ) -> Self {
         Self {
-            client,
+            client: Arc::new(client),
             peers,
             incoming_requests: ReceiverStream::new(incoming),
             header_transform,
@@ -89,20 +90,22 @@ impl<C, N: NetworkPrimitives> EthRequestHandler<C, N> {
 impl<C, N> EthRequestHandler<C, N>
 where
     N: NetworkPrimitives,
-    C: BlockReader<Header = N::BlockHeader>,
+    C: BlockReader<Header = N::BlockHeader> + 'static,
 {
     /// Returns the list of requested headers
-    fn get_headers_response(&self, request: GetBlockHeaders) -> Vec<C::Header> {
+    async fn get_headers_response(
+        client: Arc<C>,
+        header_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
+        request: GetBlockHeaders,
+    ) -> Vec<C::Header> {
         let GetBlockHeaders { start_block, limit, skip, direction } = request;
 
         let mut headers = Vec::new();
 
         let mut block: BlockHashOrNumber = match start_block {
             BlockHashOrNumber::Hash(start) => start.into(),
             BlockHashOrNumber::Number(num) => {
-                let Some(hash) = self.client.block_hash(num).unwrap_or_default() else {
-                    return headers
-                };
+                let Some(hash) = client.block_hash(num).unwrap_or_default() else { return headers };
                 hash.into()
             }
         };
@@ -111,7 +114,7 @@ where
         let mut total_bytes = 0;
 
         for _ in 0..limit {
-            if let Some(header) = self.client.header_by_hash_or_number(block).unwrap_or_default() {
+            if let Some(header) = client.header_by_hash_or_number(block).unwrap_or_default() {
                 let number = header.number();
                 let parent_hash = header.parent_hash();
 
@@ -153,8 +156,8 @@ where
         }
 
         // TODO: remove this once we deprecated l2geth
-        if let Some(ref header_transform) = self.header_transform {
-            headers = headers.into_iter().map(|h| header_transform.map(h)).collect()
+        if let Some(ref header_transform) = header_transform {
+            return join_all(headers.into_iter().map(|h| header_transform.map(h))).await;
         }
 
         headers
@@ -165,10 +168,14 @@ where
         _peer_id: PeerId,
         request: GetBlockHeaders,
         response: oneshot::Sender<RequestResult<BlockHeaders<C::Header>>>,
-    ) {
+    ) -> impl Future<Output = ()> + 'static {
         self.metrics.eth_headers_requests_received_total.increment(1);
-        let headers = self.get_headers_response(request);
-        let _ = response.send(Ok(BlockHeaders(headers)));
+        let client = self.client.clone();
+        let header_transform = self.header_transform.clone();
+        async move {
+            let headers = Self::get_headers_response(client, header_transform, request).await;
+            let _ = response.send(Ok(BlockHeaders(headers)));
+        }
     }
 
     fn on_bodies_request(
@@ -267,7 +274,8 @@ where
     N: NetworkPrimitives,
     C: BlockReader<Block = N::Block, Receipt = N::Receipt>
         + HeaderProvider<Header = N::BlockHeader>
-        + Unpin,
+        + Unpin
+        + 'static,
 {
     type Output = ();
 
@@ -284,7 +292,8 @@ where
             |incoming| {
                 match incoming {
                     IncomingEthRequest::GetBlockHeaders { peer_id, request, response } => {
-                        this.on_headers_request(peer_id, request, response)
+                        let future = this.on_headers_request(peer_id, request, response);
+                        tokio::spawn(future);
                     }
                     IncomingEthRequest::GetBlockBodies { peer_id, request, response } => {
                         this.on_bodies_request(peer_id, request, response)

crates/net/network/src/eth_requests.rs

@@ -13,3 +13,18 @@ impl<H: BlockHeader> HeaderTransform<H> for () {
         header
     }
 }
+
+/// An instance of the trait applies a mapping to headers that are being sent to a peer in response
+/// to a request.
+#[async_trait::async_trait]
+pub trait HeaderResponseTransform<H: BlockHeader>: std::fmt::Debug + Send + Sync {
+    /// Applies a mapping to the response headers.
+    async fn map(&self, header: H) -> H;
+}
+
+#[async_trait::async_trait]
+impl<H: BlockHeader> HeaderResponseTransform<H> for () {
+    async fn map(&self, header: H) -> H {
+        header
+    }
+}

crates/net/network/src/transform/header.rs

@@ -17,7 +17,7 @@ use reth_db_api::{database::Database, database_metrics::DatabaseMetrics};
 use reth_exex::ExExContext;
 use reth_network::{
     transactions::{TransactionPropagationPolicy, TransactionsManagerConfig},
-    transform::header::HeaderTransform,
+    transform::header::HeaderResponseTransform,
     NetworkBuilder, NetworkConfig, NetworkConfigBuilder, NetworkHandle, NetworkManager,
     NetworkPrimitives,
 };
@@ -586,10 +586,10 @@ where
     ///     .extend_rpc_modules(|ctx| {
     ///         // Access node components, so they can used by the CustomApi
     ///         let pool = ctx.pool().clone();
-    ///         
+    ///
     ///         // Add custom RPC namespace
     ///         ctx.modules.merge_configured(CustomApi { pool }.into_rpc())?;
-    ///         
+    ///
     ///         Ok(())
     ///     })
     ///     .build()?;
@@ -790,7 +790,7 @@ impl<Node: FullNodeTypes> BuilderContext<Node> {
         &self,
         builder: NetworkBuilder<(), (), N>,
         pool: Pool,
-        request_transform: Option<Box<dyn HeaderTransform<N::BlockHeader>>>,
+        request_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
     ) -> NetworkHandle<N>
     where
         N: NetworkPrimitives,
@@ -824,7 +824,7 @@ impl<Node: FullNodeTypes> BuilderContext<Node> {
         pool: Pool,
         tx_config: TransactionsManagerConfig,
         propagation_policy: Policy,
-        request_transform: Option<Box<dyn HeaderTransform<N::BlockHeader>>>,
+        request_transform: Option<Arc<dyn HeaderResponseTransform<N::BlockHeader>>>,
     ) -> NetworkHandle<N>
     where
         N: NetworkPrimitives,

crates/node/builder/src/builder/mod.rs

@@ -118,7 +118,7 @@ mod tests {
         ];
 
         for case in cases {
-            let receipts = vec![
+            let receipts = [
                 // 0xb0d6ee650637911394396d81172bd1c637d568ed1fbddab0daddfca399c58b53
                 OpReceipt::Deposit(OpDepositReceipt {
                         inner: Receipt {

crates/optimism/consensus/src/proof.rs

@@ -852,10 +852,12 @@ pub trait PayloadBuilder: Send + Sync + Clone {
 /// Tells the payload builder how to react to payload request if there's no payload available yet.
 ///
 /// This situation can occur if the CL requests a payload before the first payload has been built.
+#[derive(Default)]
 pub enum MissingPayloadBehaviour<Payload> {
     /// Await the regular scheduled payload process.
     AwaitInProgress,
     /// Race the in progress payload process with an empty payload.
+    #[default]
     RaceEmptyPayload,
     /// Race the in progress payload process with this job.
     RacePayload(Box<dyn FnOnce() -> Result<Payload, PayloadBuilderError> + Send>),
@@ -873,12 +875,6 @@ impl<Payload> fmt::Debug for MissingPayloadBehaviour<Payload> {
     }
 }
 
-impl<Payload> Default for MissingPayloadBehaviour<Payload> {
-    fn default() -> Self {
-        Self::RaceEmptyPayload
-    }
-}
-
 /// Checks if the new payload is better than the current best.
 ///
 /// This compares the total fees of the blocks, higher is better.

crates/payload/basic/src/lib.rs

@@ -8,22 +8,17 @@ use alloy_primitives::BlockNumber;
 #[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
 #[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
 #[cfg_attr(any(test, feature = "serde"), serde(rename_all = "lowercase"))]
+#[cfg_attr(any(test, feature = "test-utils"), derive(Default))]
 pub enum PruneMode {
     /// Prune all blocks.
+    #[cfg_attr(any(test, feature = "test-utils"), default)]
     Full,
     /// Prune blocks before the `head-N` block number. In other words, keep last N + 1 blocks.
     Distance(u64),
     /// Prune blocks before the specified block number. The specified block number is not pruned.
     Before(BlockNumber),
 }
 
-#[cfg(any(test, feature = "test-utils"))]
-impl Default for PruneMode {
-    fn default() -> Self {
-        Self::Full
-    }
-}
-
 impl PruneMode {
     /// Prune blocks up to the specified block number. The specified block number is also pruned.
     ///

crates/prune/types/src/mode.rs

@@ -8,8 +8,10 @@ use thiserror::Error;
 #[cfg_attr(any(test, feature = "reth-codec"), derive(reth_codecs::Compact))]
 #[cfg_attr(any(test, feature = "reth-codec"), reth_codecs::add_arbitrary_tests(compact))]
 #[cfg_attr(any(test, feature = "serde"), derive(serde::Serialize, serde::Deserialize))]
+#[cfg_attr(test, derive(Default))]
 pub enum PruneSegment {
     /// Prune segment responsible for the `TransactionSenders` table.
+    #[cfg_attr(test, default)]
     SenderRecovery,
     /// Prune segment responsible for the `TransactionHashNumbers` table.
     TransactionLookup,
@@ -72,10 +74,3 @@ pub enum PruneSegmentError {
     #[error("the configuration provided for {0} is invalid")]
     Configuration(PruneSegment),
 }
-
-#[cfg(test)]
-impl Default for PruneSegment {
-    fn default() -> Self {
-        Self::SenderRecovery
-    }
-}