feat: convert HeadersClient BodiesClient futures into associated types (#1063)

crates/interfaces/src/p2p/bodies/client.rs

@@ -1,21 +1,25 @@
+use std::pin::Pin;
+
 use crate::p2p::{download::DownloadClient, error::PeerRequestResult, priority::Priority};
-use async_trait::async_trait;
+use futures::Future;
 use reth_eth_wire::BlockBody;
 use reth_primitives::H256;
 
+/// The bodies future type
+pub type BodiesFut = Pin<Box<dyn Future<Output = PeerRequestResult<Vec<BlockBody>>> + Send + Sync>>;
+
 /// A client capable of downloading block bodies.
-#[async_trait]
 #[auto_impl::auto_impl(&, Arc, Box)]
 pub trait BodiesClient: DownloadClient {
+    /// The bodies type
+    type Output: Future<Output = PeerRequestResult<Vec<BlockBody>>> + Sync + Send + Unpin;
+
     /// Fetches the block body for the requested block.
-    async fn get_block_bodies(&self, hashes: Vec<H256>) -> PeerRequestResult<Vec<BlockBody>> {
-        self.get_block_bodies_with_priority(hashes, Priority::Normal).await
+    fn get_block_bodies(&self, hashes: Vec<H256>) -> Self::Output {
+        self.get_block_bodies_with_priority(hashes, Priority::Normal)
     }
 
     /// Fetches the block body for the requested block with priority
-    async fn get_block_bodies_with_priority(
-        &self,
-        hashes: Vec<H256>,
-        priority: Priority,
-    ) -> PeerRequestResult<Vec<BlockBody>>;
+    fn get_block_bodies_with_priority(&self, hashes: Vec<H256>, priority: Priority)
+        -> Self::Output;
 }

crates/interfaces/src/p2p/headers/client.rs

@@ -1,8 +1,8 @@
 use crate::p2p::{download::DownloadClient, error::PeerRequestResult, priority::Priority};
-use async_trait::async_trait;
+use futures::Future;
 pub use reth_eth_wire::BlockHeaders;
 use reth_primitives::{BlockHashOrNumber, HeadersDirection, H256, U256};
-use std::fmt::Debug;
+use std::{fmt::Debug, pin::Pin};
 
 /// The header request struct to be sent to connected peers, which
 /// will proceed to ask them to stream the requested headers to us.
@@ -16,23 +16,28 @@ pub struct HeadersRequest {
     pub direction: HeadersDirection,
 }
 
+/// The headers future type
+pub type HeadersFut = Pin<Box<dyn Future<Output = PeerRequestResult<BlockHeaders>> + Send + Sync>>;
+
 /// The block headers downloader client
-#[async_trait]
 #[auto_impl::auto_impl(&, Arc, Box)]
 pub trait HeadersClient: DownloadClient {
+    /// The headers type
+    type Output: Future<Output = PeerRequestResult<BlockHeaders>> + Sync + Send + Unpin;
+
     /// Sends the header request to the p2p network and returns the header response received from a
     /// peer.
-    async fn get_headers(&self, request: HeadersRequest) -> PeerRequestResult<BlockHeaders> {
-        self.get_headers_with_priority(request, Priority::Normal).await
+    fn get_headers(&self, request: HeadersRequest) -> Self::Output {
+        self.get_headers_with_priority(request, Priority::Normal)
     }
 
     /// Sends the header request to the p2p network with priroity set and returns the header
     /// response received from a peer.
-    async fn get_headers_with_priority(
+    fn get_headers_with_priority(
         &self,
         request: HeadersRequest,
         priority: Priority,
-    ) -> PeerRequestResult<BlockHeaders>;
+    ) -> Self::Output;
 }
 
 /// The status updater for updating the status of the p2p node

crates/interfaces/src/test_utils/bodies.rs

@@ -1,11 +1,18 @@
 use crate::p2p::{
-    bodies::client::BodiesClient, download::DownloadClient, error::PeerRequestResult,
+    bodies::client::{BodiesClient, BodiesFut},
+    download::DownloadClient,
+    error::PeerRequestResult,
     priority::Priority,
 };
 use async_trait::async_trait;
+use futures::{future, Future, FutureExt};
 use reth_eth_wire::BlockBody;
-use reth_primitives::H256;
-use std::fmt::{Debug, Formatter};
+use reth_primitives::{WithPeerId, H256};
+use std::{
+    fmt::{Debug, Formatter},
+    pin::Pin,
+};
+use tokio::sync::oneshot::{self, Receiver};
 
 /// A test client for fetching bodies
 pub struct TestBodiesClient<F> {
@@ -29,16 +36,22 @@ impl<F: Sync + Send> DownloadClient for TestBodiesClient<F> {
     }
 }
 
-#[async_trait]
 impl<F> BodiesClient for TestBodiesClient<F>
 where
     F: Fn(Vec<H256>) -> PeerRequestResult<Vec<BlockBody>> + Send + Sync,
 {
-    async fn get_block_bodies_with_priority(
+    type Output = BodiesFut;
+
+    fn get_block_bodies_with_priority(
         &self,
         hashes: Vec<H256>,
         _priority: Priority,
-    ) -> PeerRequestResult<Vec<BlockBody>> {
-        (self.responder)(hashes)
+    ) -> Self::Output {
+        let (tx, rx) = oneshot::channel();
+        tx.send((self.responder)(hashes));
+        Box::pin(rx.map(|x| match x {
+            Ok(value) => value,
+            Err(err) => Err(err.into()),
+        }))
     }
 }

crates/interfaces/src/test_utils/headers.rs

@@ -11,10 +11,11 @@ use crate::{
         priority::Priority,
     },
 };
-use futures::{Future, FutureExt, Stream, StreamExt};
+use futures::{future, Future, FutureExt, Stream, StreamExt};
 use reth_eth_wire::BlockHeaders;
 use reth_primitives::{
-    BlockHash, BlockNumber, Header, HeadersDirection, PeerId, SealedBlock, SealedHeader, H256,
+    BlockHash, BlockNumber, Header, HeadersDirection, PeerId, SealedBlock, SealedHeader,
+    WithPeerId, H256,
 };
 use reth_rpc_types::engine::ForkchoiceState;
 use std::{
@@ -26,7 +27,12 @@ use std::{
     },
     task::{ready, Context, Poll},
 };
-use tokio::sync::{watch, watch::error::SendError, Mutex};
+use tokio::sync::{
+    oneshot::{error::RecvError, Receiver},
+    watch,
+    watch::error::SendError,
+    Mutex,
+};
 
 /// A test downloader which just returns the values that have been pushed to it.
 #[derive(Debug)]
@@ -98,7 +104,7 @@ impl Stream for TestHeaderDownloader {
     }
 }
 
-type TestHeadersFut = Pin<Box<dyn Future<Output = PeerRequestResult<BlockHeaders>> + Send>>;
+type TestHeadersFut = Pin<Box<dyn Future<Output = PeerRequestResult<BlockHeaders>> + Sync + Send>>;
 
 struct TestDownload {
     client: Arc<TestHeadersClient>,
@@ -109,9 +115,6 @@ struct TestDownload {
     done: bool,
 }
 
-/// SAFETY: All the mutations are performed through an exclusive reference on `poll`
-unsafe impl Sync for TestDownload {}
-
 impl TestDownload {
     fn get_or_init_fut(&mut self) -> &mut TestHeadersFut {
         if self.fut.is_none() {
@@ -121,7 +124,7 @@ impl TestDownload {
                 start: reth_primitives::BlockHashOrNumber::Number(0), // ignored
             };
             let client = Arc::clone(&self.client);
-            self.fut = Some(Box::pin(async move { client.get_headers(request).await }));
+            self.fut = Some(Box::pin(client.get_headers(request)));
         }
         self.fut.as_mut().unwrap()
     }
@@ -226,22 +229,30 @@ impl DownloadClient for TestHeadersClient {
     }
 }
 
-#[async_trait::async_trait]
 impl HeadersClient for TestHeadersClient {
-    async fn get_headers_with_priority(
+    type Output = TestHeadersFut;
+
+    fn get_headers_with_priority(
         &self,
         request: HeadersRequest,
         _priority: Priority,
-    ) -> PeerRequestResult<BlockHeaders> {
-        self.request_attempts.fetch_add(1, Ordering::SeqCst);
-        if let Some(err) = &mut *self.error.lock().await {
-            return Err(err.clone())
-        }
+    ) -> Self::Output {
+        let responses = self.responses.clone();
+        let error = self.error.clone();
 
-        let mut lock = self.responses.lock().await;
-        let len = lock.len().min(request.limit as usize);
-        let resp = lock.drain(..len).collect();
-        return Ok((PeerId::default(), BlockHeaders(resp)).into())
+        self.request_attempts.fetch_add(1, Ordering::SeqCst);
+
+        Box::pin(async move {
+            if let Some(err) = &mut *error.lock().await {
+                return Err(err.clone())
+            }
+
+            let mut lock = responses.lock().await;
+            let len = lock.len().min(request.limit as usize);
+            let resp = lock.drain(..len).collect();
+            let with_peer_id = WithPeerId::from((PeerId::default(), BlockHeaders(resp)));
+            Ok(with_peer_id)
+        })
     }
 }