Skip to content

Commit

Permalink
feat: make BlockResponse generic over header (#13195)
Browse files Browse the repository at this point in the history
  • Loading branch information
Rjected authored Dec 6, 2024
1 parent e991570 commit 552c623
Show file tree
Hide file tree
Showing 8 changed files with 33 additions and 23 deletions.
19 changes: 13 additions & 6 deletions crates/net/downloaders/src/bodies/bodies.rs
Original file line number Diff line number Diff line change
Expand Up @@ -61,7 +61,7 @@ pub struct BodiesDownloader<B: BodiesClient, Provider> {
/// Buffered responses
buffered_responses: BinaryHeap<OrderedBodiesResponse<B::Body>>,
/// Queued body responses that can be returned for insertion into the database.
queued_bodies: Vec<BlockResponse<B::Body>>,
queued_bodies: Vec<BlockResponse<alloy_consensus::Header, B::Body>>,
/// The bodies downloader metrics.
metrics: BodyDownloaderMetrics,
}
Expand Down Expand Up @@ -193,7 +193,7 @@ where
}

/// Queues bodies and sets the latest queued block number
fn queue_bodies(&mut self, bodies: Vec<BlockResponse<B::Body>>) {
fn queue_bodies(&mut self, bodies: Vec<BlockResponse<alloy_consensus::Header, B::Body>>) {
self.latest_queued_block_number = Some(bodies.last().expect("is not empty").block_number());
self.queued_bodies.extend(bodies);
self.metrics.queued_blocks.set(self.queued_bodies.len() as f64);
Expand All @@ -210,7 +210,10 @@ where
}

/// Adds a new response to the internal buffer
fn buffer_bodies_response(&mut self, response: Vec<BlockResponse<B::Body>>) {
fn buffer_bodies_response(
&mut self,
response: Vec<BlockResponse<alloy_consensus::Header, B::Body>>,
) {
// take into account capacity
let size = response.iter().map(BlockResponse::size).sum::<usize>() +
response.capacity() * mem::size_of::<BlockResponse<B::Body>>();
Expand All @@ -227,7 +230,9 @@ where
}

/// Returns a response if it's first block number matches the next expected.
fn try_next_buffered(&mut self) -> Option<Vec<BlockResponse<B::Body>>> {
fn try_next_buffered(
&mut self,
) -> Option<Vec<BlockResponse<alloy_consensus::Header, B::Body>>> {
if let Some(next) = self.buffered_responses.peek() {
let expected = self.next_expected_block_number();
let next_block_range = next.block_range();
Expand All @@ -253,7 +258,9 @@ where

/// Returns the next batch of block bodies that can be returned if we have enough buffered
/// bodies
fn try_split_next_batch(&mut self) -> Option<Vec<BlockResponse<B::Body>>> {
fn try_split_next_batch(
&mut self,
) -> Option<Vec<BlockResponse<alloy_consensus::Header, B::Body>>> {
if self.queued_bodies.len() >= self.stream_batch_size {
let next_batch = self.queued_bodies.drain(..self.stream_batch_size).collect::<Vec<_>>();
self.queued_bodies.shrink_to_fit();
Expand Down Expand Up @@ -436,7 +443,7 @@ where

#[derive(Debug)]
struct OrderedBodiesResponse<B> {
resp: Vec<BlockResponse<B>>,
resp: Vec<BlockResponse<alloy_consensus::Header, B>>,
/// The total size of the response in bytes
size: usize,
}
Expand Down
2 changes: 1 addition & 1 deletion crates/net/downloaders/src/bodies/noop.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ impl BodyDownloader for NoopBodiesDownloader {
}

impl Stream for NoopBodiesDownloader {
type Item = Result<Vec<BlockResponse<BlockBody>>, DownloadError>;
type Item = Result<Vec<BlockResponse<alloy_consensus::Header, BlockBody>>, DownloadError>;

fn poll_next(
self: std::pin::Pin<&mut Self>,
Expand Down
2 changes: 1 addition & 1 deletion crates/net/downloaders/src/bodies/queue.rs
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ impl<B> Stream for BodiesRequestQueue<B>
where
B: BodiesClient<Body: InMemorySize> + 'static,
{
type Item = DownloadResult<Vec<BlockResponse<B::Body>>>;
type Item = DownloadResult<Vec<BlockResponse<alloy_consensus::Header, B::Body>>>;

fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.get_mut().inner.poll_next_unpin(cx)
Expand Down
4 changes: 2 additions & 2 deletions crates/net/downloaders/src/bodies/request.rs
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ pub(crate) struct BodiesRequestFuture<B: BodiesClient> {
// Headers to download. The collection is shrunk as responses are buffered.
pending_headers: VecDeque<SealedHeader>,
/// Internal buffer for all blocks
buffer: Vec<BlockResponse<B::Body>>,
buffer: Vec<BlockResponse<alloy_consensus::Header, B::Body>>,
fut: Option<B::Output>,
/// Tracks how many bodies we requested in the last request.
last_request_len: Option<usize>,
Expand Down Expand Up @@ -217,7 +217,7 @@ impl<B> Future for BodiesRequestFuture<B>
where
B: BodiesClient<Body: InMemorySize> + 'static,
{
type Output = DownloadResult<Vec<BlockResponse<B::Body>>>;
type Output = DownloadResult<Vec<BlockResponse<alloy_consensus::Header, B::Body>>>;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
Expand Down
2 changes: 1 addition & 1 deletion crates/net/p2p/src/bodies/downloader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use futures::Stream;
use std::{fmt::Debug, ops::RangeInclusive};

/// Body downloader return type.
pub type BodyDownloaderResult<B> = DownloadResult<Vec<BlockResponse<B>>>;
pub type BodyDownloaderResult<B> = DownloadResult<Vec<BlockResponse<alloy_consensus::Header, B>>>;

/// A downloader capable of fetching and yielding block bodies from block headers.
///
Expand Down
23 changes: 13 additions & 10 deletions crates/net/p2p/src/bodies/response.rs
Original file line number Diff line number Diff line change
@@ -1,19 +1,22 @@
use alloy_primitives::{BlockNumber, U256};
use reth_primitives::{BlockBody, SealedBlock, SealedHeader};
use reth_primitives_traits::InMemorySize;
use reth_primitives_traits::{BlockHeader, InMemorySize};

/// The block response
#[derive(PartialEq, Eq, Debug, Clone)]
pub enum BlockResponse<B = BlockBody> {
pub enum BlockResponse<H = alloy_consensus::Header, B = BlockBody> {
/// Full block response (with transactions or ommers)
Full(SealedBlock<alloy_consensus::Header, B>),
Full(SealedBlock<H, B>),
/// The empty block response
Empty(SealedHeader),
Empty(SealedHeader<H>),
}

impl<B> BlockResponse<B> {
impl<H, B> BlockResponse<H, B>
where
H: BlockHeader,
{
/// Return the reference to the response header
pub const fn header(&self) -> &SealedHeader {
pub const fn header(&self) -> &SealedHeader<H> {
match self {
Self::Full(block) => &block.header,
Self::Empty(header) => header,
Expand All @@ -22,14 +25,14 @@ impl<B> BlockResponse<B> {

/// Return the block number
pub fn block_number(&self) -> BlockNumber {
self.header().number
self.header().number()
}

/// Return the reference to the response header
pub fn difficulty(&self) -> U256 {
match self {
Self::Full(block) => block.difficulty,
Self::Empty(header) => header.difficulty,
Self::Full(block) => block.difficulty(),
Self::Empty(header) => header.difficulty(),
}
}

Expand All @@ -42,7 +45,7 @@ impl<B> BlockResponse<B> {
}
}

impl<B: InMemorySize> InMemorySize for BlockResponse<B> {
impl<H: InMemorySize, B: InMemorySize> InMemorySize for BlockResponse<H, B> {
#[inline]
fn size(&self) -> usize {
match self {
Expand Down
2 changes: 1 addition & 1 deletion crates/primitives-traits/src/header/sealed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ impl<H: alloy_consensus::BlockHeader> SealedHeader<H> {
}
}

impl InMemorySize for SealedHeader {
impl<H: InMemorySize> InMemorySize for SealedHeader<H> {
/// Calculates a heuristic for the in-memory size of the [`SealedHeader`].
#[inline]
fn size(&self) -> usize {
Expand Down
2 changes: 1 addition & 1 deletion crates/stages/stages/src/stages/bodies.rs
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ pub struct BodyStage<D: BodyDownloader> {
/// The body downloader.
downloader: D,
/// Block response buffer.
buffer: Option<Vec<BlockResponse<D::Body>>>,
buffer: Option<Vec<BlockResponse<alloy_consensus::Header, D::Body>>>,
}

impl<D: BodyDownloader> BodyStage<D> {
Expand Down

0 comments on commit 552c623

Please sign in to comment.