reth_downloaders/bodies/
bodies.rs

1use super::queue::BodiesRequestQueue;
2use crate::{bodies::task::TaskDownloader, metrics::BodyDownloaderMetrics};
3use alloy_consensus::BlockHeader;
4use alloy_primitives::BlockNumber;
5use futures::Stream;
6use futures_util::StreamExt;
7use reth_config::BodiesConfig;
8use reth_consensus::{Consensus, ConsensusError};
9use reth_network_p2p::{
10    bodies::{
11        client::BodiesClient,
12        downloader::{BodyDownloader, BodyDownloaderResult},
13        response::BlockResponse,
14    },
15    error::{DownloadError, DownloadResult},
16};
17use reth_primitives_traits::{size::InMemorySize, Block, SealedHeader};
18use reth_storage_api::HeaderProvider;
19use reth_tasks::{TaskSpawner, TokioTaskExecutor};
20use std::{
21    cmp::Ordering,
22    collections::BinaryHeap,
23    fmt::Debug,
24    mem,
25    ops::RangeInclusive,
26    pin::Pin,
27    sync::Arc,
28    task::{Context, Poll},
29};
30use tracing::info;
31
32/// Downloads bodies in batches.
33///
34/// All blocks in a batch are fetched at the same time.
35#[must_use = "Stream does nothing unless polled"]
36#[derive(Debug)]
37pub struct BodiesDownloader<
38    B: Block,
39    C: BodiesClient<Body = B::Body>,
40    Provider: HeaderProvider<Header = B::Header>,
41> {
42    /// The bodies client
43    client: Arc<C>,
44    /// The consensus client
45    consensus: Arc<dyn Consensus<B, Error = ConsensusError>>,
46    /// The database handle
47    provider: Provider,
48    /// The maximum number of non-empty blocks per one request
49    request_limit: u64,
50    /// The maximum number of block bodies returned at once from the stream
51    stream_batch_size: usize,
52    /// The allowed range for number of concurrent requests.
53    concurrent_requests_range: RangeInclusive<usize>,
54    /// Maximum number of bytes of received blocks to buffer internally.
55    max_buffered_blocks_size_bytes: usize,
56    /// Current estimated size of buffered blocks in bytes.
57    buffered_blocks_size_bytes: usize,
58    /// The range of block numbers for body download.
59    download_range: RangeInclusive<BlockNumber>,
60    /// The latest block number returned.
61    latest_queued_block_number: Option<BlockNumber>,
62    /// Requests in progress
63    in_progress_queue: BodiesRequestQueue<B, C>,
64    /// Buffered responses
65    buffered_responses: BinaryHeap<OrderedBodiesResponse<B>>,
66    /// Queued body responses that can be returned for insertion into the database.
67    queued_bodies: Vec<BlockResponse<B>>,
68    /// The bodies downloader metrics.
69    metrics: BodyDownloaderMetrics,
70}
71
72impl<B, C, Provider> BodiesDownloader<B, C, Provider>
73where
74    B: Block,
75    C: BodiesClient<Body = B::Body> + 'static,
76    Provider: HeaderProvider<Header = B::Header> + Unpin + 'static,
77{
78    /// Returns the next contiguous request.
79    fn next_headers_request(&self) -> DownloadResult<Option<Vec<SealedHeader<Provider::Header>>>> {
80        let start_at = match self.in_progress_queue.last_requested_block_number {
81            Some(num) => num + 1,
82            None => *self.download_range.start(),
83        };
84        // as the range is inclusive, we need to add 1 to the end.
85        let items_left = (self.download_range.end() + 1).saturating_sub(start_at);
86        let limit = items_left.min(self.request_limit);
87        self.query_headers(start_at..=*self.download_range.end(), limit)
88    }
89
90    /// Retrieve a batch of headers from the database starting from the provided block number.
91    ///
92    /// This method is going to return the batch as soon as one of the conditions below
93    /// is fulfilled:
94    ///     1. The number of non-empty headers in the batch equals requested.
95    ///     2. The total number of headers in the batch (both empty and non-empty) is greater than
96    ///        or equal to the stream batch size.
97    ///     3. Downloader reached the end of the range
98    ///
99    /// NOTE: The batches returned have a variable length.
100    fn query_headers(
101        &self,
102        range: RangeInclusive<BlockNumber>,
103        max_non_empty: u64,
104    ) -> DownloadResult<Option<Vec<SealedHeader<B::Header>>>> {
105        if range.is_empty() || max_non_empty == 0 {
106            return Ok(None)
107        }
108
109        // Collect headers while
110        //      1. Current block number is in range
111        //      2. The number of non empty headers is less than maximum
112        //      3. The total number of headers is less than the stream batch size (this is only
113        //         relevant if the range consists entirely of empty headers)
114        let mut collected = 0;
115        let mut non_empty_headers = 0;
116        let headers = self.provider.sealed_headers_while(range.clone(), |header| {
117            let should_take = range.contains(&header.number()) &&
118                non_empty_headers < max_non_empty &&
119                collected < self.stream_batch_size;
120
121            if should_take {
122                collected += 1;
123                if !header.is_empty() {
124                    non_empty_headers += 1;
125                }
126                true
127            } else {
128                false
129            }
130        })?;
131
132        Ok(Some(headers).filter(|h| !h.is_empty()))
133    }
134
135    /// Get the next expected block number for queueing.
136    const fn next_expected_block_number(&self) -> BlockNumber {
137        match self.latest_queued_block_number {
138            Some(num) => num + 1,
139            None => *self.download_range.start(),
140        }
141    }
142
143    /// Max requests to handle at the same time
144    ///
145    /// This depends on the number of active peers but will always be
146    /// `min_concurrent_requests..max_concurrent_requests`
147    #[inline]
148    fn concurrent_request_limit(&self) -> usize {
149        let num_peers = self.client.num_connected_peers();
150
151        let max_requests = num_peers.max(*self.concurrent_requests_range.start());
152
153        // if we're only connected to a few peers, we keep it low
154        if num_peers < *self.concurrent_requests_range.start() {
155            return max_requests
156        }
157
158        max_requests.min(*self.concurrent_requests_range.end())
159    }
160
161    /// Returns true if the size of buffered blocks is lower than the configured maximum
162    const fn has_buffer_capacity(&self) -> bool {
163        self.buffered_blocks_size_bytes < self.max_buffered_blocks_size_bytes
164    }
165
166    // Check if the stream is terminated
167    fn is_terminated(&self) -> bool {
168        // There is nothing to request if the range is empty
169        let nothing_to_request = self.download_range.is_empty() ||
170            // or all blocks have already been requested.
171            self.in_progress_queue
172                .last_requested_block_number.is_some_and(|last| last == *self.download_range.end());
173
174        nothing_to_request &&
175            self.in_progress_queue.is_empty() &&
176            self.buffered_responses.is_empty() &&
177            self.queued_bodies.is_empty()
178    }
179
180    /// Clear all download related data.
181    ///
182    /// Should be invoked upon encountering fatal error.
183    fn clear(&mut self) {
184        self.download_range = RangeInclusive::new(1, 0);
185        self.latest_queued_block_number.take();
186        self.in_progress_queue.clear();
187        self.queued_bodies = Vec::new();
188        self.buffered_responses = BinaryHeap::new();
189        self.buffered_blocks_size_bytes = 0;
190
191        // reset metrics
192        self.metrics.in_flight_requests.set(0.);
193        self.metrics.buffered_responses.set(0.);
194        self.metrics.buffered_blocks.set(0.);
195        self.metrics.buffered_blocks_size_bytes.set(0.);
196        self.metrics.queued_blocks.set(0.);
197    }
198
199    /// Queues bodies and sets the latest queued block number
200    fn queue_bodies(&mut self, bodies: Vec<BlockResponse<B>>) {
201        self.latest_queued_block_number = Some(bodies.last().expect("is not empty").block_number());
202        self.queued_bodies.extend(bodies);
203        self.metrics.queued_blocks.set(self.queued_bodies.len() as f64);
204    }
205
206    /// Removes the next response from the buffer.
207    fn pop_buffered_response(&mut self) -> Option<OrderedBodiesResponse<B>> {
208        let resp = self.buffered_responses.pop()?;
209        self.metrics.buffered_responses.decrement(1.);
210        self.buffered_blocks_size_bytes -= resp.size();
211        self.metrics.buffered_blocks.decrement(resp.len() as f64);
212        self.metrics.buffered_blocks_size_bytes.set(self.buffered_blocks_size_bytes as f64);
213        Some(resp)
214    }
215
216    /// Adds a new response to the internal buffer
217    fn buffer_bodies_response(&mut self, response: Vec<BlockResponse<B>>) {
218        // take into account capacity
219        let size = response.iter().map(BlockResponse::size).sum::<usize>() +
220            response.capacity() * mem::size_of::<BlockResponse<B>>();
221
222        let response = OrderedBodiesResponse { resp: response, size };
223        let response_len = response.len();
224
225        self.buffered_blocks_size_bytes += size;
226        self.buffered_responses.push(response);
227
228        self.metrics.buffered_blocks.increment(response_len as f64);
229        self.metrics.buffered_blocks_size_bytes.set(self.buffered_blocks_size_bytes as f64);
230        self.metrics.buffered_responses.set(self.buffered_responses.len() as f64);
231    }
232
233    /// Returns a response if its first block number matches the next expected.
234    fn try_next_buffered(&mut self) -> Option<Vec<BlockResponse<B>>> {
235        if let Some(next) = self.buffered_responses.peek() {
236            let expected = self.next_expected_block_number();
237            let next_block_range = next.block_range();
238
239            if next_block_range.contains(&expected) {
240                return self.pop_buffered_response().map(|buffered| {
241                    buffered
242                        .resp
243                        .into_iter()
244                        .skip_while(|b| b.block_number() < expected)
245                        .take_while(|b| self.download_range.contains(&b.block_number()))
246                        .collect()
247                })
248            }
249
250            // Drop buffered response since we passed that range
251            if *next_block_range.end() < expected {
252                self.pop_buffered_response();
253            }
254        }
255        None
256    }
257
258    /// Returns the next batch of block bodies that can be returned if we have enough buffered
259    /// bodies
260    fn try_split_next_batch(&mut self) -> Option<Vec<BlockResponse<B>>> {
261        if self.queued_bodies.len() >= self.stream_batch_size {
262            let next_batch = self.queued_bodies.drain(..self.stream_batch_size).collect::<Vec<_>>();
263            self.queued_bodies.shrink_to_fit();
264            self.metrics.total_flushed.increment(next_batch.len() as u64);
265            self.metrics.queued_blocks.set(self.queued_bodies.len() as f64);
266            return Some(next_batch)
267        }
268        None
269    }
270
271    /// Check if a new request can be submitted, it implements back pressure to prevent overwhelming
272    /// the system and causing memory overload.
273    ///
274    /// Returns true if a new request can be submitted
275    fn can_submit_new_request(&self) -> bool {
276        // requests are issued in order but not necessarily finished in order, so the queued bodies
277        // can grow large if a certain request is slow, so we limit the followup requests if the
278        // queued bodies grew too large
279        self.queued_bodies.len() < 4 * self.stream_batch_size &&
280            self.has_buffer_capacity() &&
281            self.in_progress_queue.len() < self.concurrent_request_limit()
282    }
283}
284
285impl<B, C, Provider> BodiesDownloader<B, C, Provider>
286where
287    B: Block + 'static,
288    C: BodiesClient<Body = B::Body> + 'static,
289    Provider: HeaderProvider<Header = B::Header> + Unpin + 'static,
290{
291    /// Spawns the downloader task via [`tokio::task::spawn`]
292    pub fn into_task(self) -> TaskDownloader<B> {
293        self.into_task_with(&TokioTaskExecutor::default())
294    }
295
296    /// Convert the downloader into a [`TaskDownloader`] by spawning it via the given spawner.
297    pub fn into_task_with<S>(self, spawner: &S) -> TaskDownloader<B>
298    where
299        S: TaskSpawner,
300    {
301        TaskDownloader::spawn_with(self, spawner)
302    }
303}
304
305impl<B, C, Provider> BodyDownloader for BodiesDownloader<B, C, Provider>
306where
307    B: Block + 'static,
308    C: BodiesClient<Body = B::Body> + 'static,
309    Provider: HeaderProvider<Header = B::Header> + Unpin + 'static,
310{
311    type Block = B;
312
313    /// Set a new download range (exclusive).
314    ///
315    /// This method will drain all queued bodies, filter out ones outside the range and put them
316    /// back into the buffer.
317    /// If there are any bodies between the range start and last queued body that have not been
318    /// downloaded or are not in progress, they will be re-requested.
319    fn set_download_range(&mut self, range: RangeInclusive<BlockNumber>) -> DownloadResult<()> {
320        // Check if the range is valid.
321        if range.is_empty() {
322            tracing::error!(target: "downloaders::bodies", ?range, "Bodies download range is invalid (empty)");
323            return Err(DownloadError::InvalidBodyRange { range })
324        }
325
326        // Check if the provided range is the subset of the existing range.
327        let is_current_range_subset = self.download_range.contains(range.start()) &&
328            *range.end() == *self.download_range.end();
329        if is_current_range_subset {
330            tracing::trace!(target: "downloaders::bodies", ?range, "Download range already in progress");
331            // The current range already includes requested.
332            return Ok(())
333        }
334
335        // Check if the provided range is the next expected range.
336        let count = *range.end() - *range.start() + 1; // range is inclusive
337        let is_next_consecutive_range = *range.start() == *self.download_range.end() + 1;
338        if is_next_consecutive_range {
339            // New range received.
340            tracing::trace!(target: "downloaders::bodies", ?range, "New download range set");
341            info!(target: "downloaders::bodies", count, ?range, "Downloading bodies");
342            self.download_range = range;
343            return Ok(())
344        }
345
346        // The block range is reset. This can happen either after unwind or after the bodies were
347        // written by external services (e.g. BlockchainTree).
348        tracing::trace!(target: "downloaders::bodies", ?range, prev_range = ?self.download_range, "Download range reset");
349        info!(target: "downloaders::bodies", count, ?range, "Downloading bodies");
350        // Increment out-of-order requests metric if the new start is below the last returned block
351        if let Some(last_returned) = self.latest_queued_block_number &&
352            *range.start() < last_returned
353        {
354            self.metrics.out_of_order_requests.increment(1);
355        }
356        self.clear();
357        self.download_range = range;
358        Ok(())
359    }
360}
361
362impl<B, C, Provider> Stream for BodiesDownloader<B, C, Provider>
363where
364    B: Block + 'static,
365    C: BodiesClient<Body = B::Body> + 'static,
366    Provider: HeaderProvider<Header = B::Header> + Unpin + 'static,
367{
368    type Item = BodyDownloaderResult<B>;
369
370    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
371        let this = self.get_mut();
372        if this.is_terminated() {
373            return Poll::Ready(None)
374        }
375        // Submit new requests and poll any in progress
376        loop {
377            // Yield next batch if ready
378            if let Some(next_batch) = this.try_split_next_batch() {
379                return Poll::Ready(Some(Ok(next_batch)))
380            }
381
382            // Poll requests
383            while let Poll::Ready(Some(response)) = this.in_progress_queue.poll_next_unpin(cx) {
384                this.metrics.in_flight_requests.decrement(1.);
385                match response {
386                    Ok(response) => {
387                        this.buffer_bodies_response(response);
388                    }
389                    Err(error) => {
390                        tracing::debug!(target: "downloaders::bodies", %error, "Request failed");
391                        this.clear();
392                        return Poll::Ready(Some(Err(error)))
393                    }
394                };
395            }
396
397            // Loop exit condition
398            let mut new_request_submitted = false;
399            // Submit new requests
400            'inner: while this.can_submit_new_request() {
401                match this.next_headers_request() {
402                    Ok(Some(request)) => {
403                        this.metrics.in_flight_requests.increment(1.);
404                        this.in_progress_queue.push_new_request(
405                            Arc::clone(&this.client),
406                            Arc::clone(&this.consensus),
407                            request,
408                        );
409                        new_request_submitted = true;
410                    }
411                    Ok(None) => break 'inner,
412                    Err(error) => {
413                        tracing::error!(target: "downloaders::bodies", %error, "Failed to download from next request");
414                        this.clear();
415                        return Poll::Ready(Some(Err(error)))
416                    }
417                };
418            }
419
420            while let Some(buf_response) = this.try_next_buffered() {
421                this.queue_bodies(buf_response);
422            }
423
424            // shrink the buffer so that it doesn't grow indefinitely
425            this.buffered_responses.shrink_to_fit();
426
427            if !new_request_submitted {
428                break
429            }
430        }
431
432        // All requests are handled, stream is finished
433        if this.in_progress_queue.is_empty() {
434            if this.queued_bodies.is_empty() {
435                return Poll::Ready(None)
436            }
437            let batch_size = this.stream_batch_size.min(this.queued_bodies.len());
438            let next_batch = this.queued_bodies.drain(..batch_size).collect::<Vec<_>>();
439            this.queued_bodies.shrink_to_fit();
440            this.metrics.total_flushed.increment(next_batch.len() as u64);
441            this.metrics.queued_blocks.set(this.queued_bodies.len() as f64);
442            return Poll::Ready(Some(Ok(next_batch)))
443        }
444
445        Poll::Pending
446    }
447}
448
449#[derive(Debug)]
450struct OrderedBodiesResponse<B: Block> {
451    resp: Vec<BlockResponse<B>>,
452    /// The total size of the response in bytes
453    size: usize,
454}
455
456impl<B: Block> OrderedBodiesResponse<B> {
457    #[inline]
458    const fn len(&self) -> usize {
459        self.resp.len()
460    }
461
462    /// Returns the size of the response in bytes
463    ///
464    /// See [`BlockResponse::size`]
465    #[inline]
466    const fn size(&self) -> usize {
467        self.size
468    }
469}
470
471impl<B: Block> OrderedBodiesResponse<B> {
472    /// Returns the block number of the first element
473    ///
474    /// # Panics
475    /// If the response vec is empty.
476    fn first_block_number(&self) -> u64 {
477        self.resp.first().expect("is not empty").block_number()
478    }
479
480    /// Returns the range of the block numbers in the response
481    ///
482    /// # Panics
483    /// If the response vec is empty.
484    fn block_range(&self) -> RangeInclusive<u64> {
485        self.first_block_number()..=self.resp.last().expect("is not empty").block_number()
486    }
487}
488
489impl<B: Block> PartialEq for OrderedBodiesResponse<B> {
490    fn eq(&self, other: &Self) -> bool {
491        self.first_block_number() == other.first_block_number()
492    }
493}
494
495impl<B: Block> Eq for OrderedBodiesResponse<B> {}
496
497impl<B: Block> PartialOrd for OrderedBodiesResponse<B> {
498    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
499        Some(self.cmp(other))
500    }
501}
502
503impl<B: Block> Ord for OrderedBodiesResponse<B> {
504    fn cmp(&self, other: &Self) -> Ordering {
505        self.first_block_number().cmp(&other.first_block_number()).reverse()
506    }
507}
508
509/// Builder for [`BodiesDownloader`].
510#[derive(Debug, Clone)]
511pub struct BodiesDownloaderBuilder {
512    /// The batch size of non-empty blocks per one request
513    pub request_limit: u64,
514    /// The maximum number of block bodies returned at once from the stream
515    pub stream_batch_size: usize,
516    /// Maximum number of bytes of received bodies to buffer internally.
517    pub max_buffered_blocks_size_bytes: usize,
518    /// The maximum number of requests to send concurrently.
519    pub concurrent_requests_range: RangeInclusive<usize>,
520}
521
522impl BodiesDownloaderBuilder {
523    /// Creates a new [`BodiesDownloaderBuilder`] with configurations based on the provided
524    /// [`BodiesConfig`].
525    pub fn new(config: BodiesConfig) -> Self {
526        Self::default()
527            .with_stream_batch_size(config.downloader_stream_batch_size)
528            .with_request_limit(config.downloader_request_limit)
529            .with_max_buffered_blocks_size_bytes(config.downloader_max_buffered_blocks_size_bytes)
530            .with_concurrent_requests_range(
531                config.downloader_min_concurrent_requests..=
532                    config.downloader_max_concurrent_requests,
533            )
534    }
535}
536
537impl Default for BodiesDownloaderBuilder {
538    fn default() -> Self {
539        Self {
540            request_limit: 200,
541            stream_batch_size: 1_000,
542            max_buffered_blocks_size_bytes: 2 * 1024 * 1024 * 1024, // ~2GB
543            concurrent_requests_range: 5..=100,
544        }
545    }
546}
547
548impl BodiesDownloaderBuilder {
549    /// Set request batch size on the downloader.
550    pub const fn with_request_limit(mut self, request_limit: u64) -> Self {
551        self.request_limit = request_limit;
552        self
553    }
554
555    /// Set stream batch size on the downloader.
556    pub const fn with_stream_batch_size(mut self, stream_batch_size: usize) -> Self {
557        self.stream_batch_size = stream_batch_size;
558        self
559    }
560
561    /// Set concurrent requests range on the downloader.
562    pub const fn with_concurrent_requests_range(
563        mut self,
564        concurrent_requests_range: RangeInclusive<usize>,
565    ) -> Self {
566        self.concurrent_requests_range = concurrent_requests_range;
567        self
568    }
569
570    /// Set max buffered block bytes on the downloader.
571    pub const fn with_max_buffered_blocks_size_bytes(
572        mut self,
573        max_buffered_blocks_size_bytes: usize,
574    ) -> Self {
575        self.max_buffered_blocks_size_bytes = max_buffered_blocks_size_bytes;
576        self
577    }
578
579    /// Consume self and return the concurrent downloader.
580    pub fn build<B, C, Provider>(
581        self,
582        client: C,
583        consensus: Arc<dyn Consensus<B, Error = ConsensusError>>,
584        provider: Provider,
585    ) -> BodiesDownloader<B, C, Provider>
586    where
587        B: Block,
588        C: BodiesClient<Body = B::Body> + 'static,
589        Provider: HeaderProvider<Header = B::Header>,
590    {
591        let Self {
592            request_limit,
593            stream_batch_size,
594            concurrent_requests_range,
595            max_buffered_blocks_size_bytes,
596        } = self;
597        let metrics = BodyDownloaderMetrics::default();
598        let in_progress_queue = BodiesRequestQueue::new(metrics.clone());
599        BodiesDownloader {
600            client: Arc::new(client),
601            consensus,
602            provider,
603            request_limit,
604            stream_batch_size,
605            max_buffered_blocks_size_bytes,
606            concurrent_requests_range,
607            in_progress_queue,
608            metrics,
609            download_range: RangeInclusive::new(1, 0),
610            latest_queued_block_number: None,
611            buffered_responses: Default::default(),
612            queued_bodies: Default::default(),
613            buffered_blocks_size_bytes: 0,
614        }
615    }
616}
617
618#[cfg(test)]
619mod tests {
620    use super::*;
621    use crate::{
622        bodies::test_utils::{insert_headers, zip_blocks},
623        test_utils::{generate_bodies, TestBodiesClient},
624    };
625    use alloy_primitives::B256;
626    use assert_matches::assert_matches;
627    use reth_consensus::test_utils::TestConsensus;
628    use reth_provider::test_utils::create_test_provider_factory;
629    use reth_testing_utils::generators::{self, random_block_range, BlockRangeParams};
630    use std::collections::HashMap;
631
632    // Check that the blocks are emitted in order of block number, not in order of
633    // first-downloaded
634    #[tokio::test]
635    async fn streams_bodies_in_order() {
636        // Generate some random blocks
637        let factory = create_test_provider_factory();
638        let (headers, mut bodies) = generate_bodies(0..=19);
639
640        insert_headers(&factory, &headers);
641
642        let client = Arc::new(
643            TestBodiesClient::default().with_bodies(bodies.clone()).with_should_delay(true),
644        );
645
646        let mut downloader = BodiesDownloaderBuilder::default()
647            .build::<reth_ethereum_primitives::Block, _, _>(
648                client.clone(),
649                Arc::new(TestConsensus::default()),
650                factory,
651            );
652        downloader.set_download_range(0..=19).expect("failed to set download range");
653
654        assert_matches!(
655            downloader.next().await,
656            Some(Ok(res)) => assert_eq!(res, zip_blocks(headers.iter(), &mut bodies))
657        );
658        assert_eq!(client.times_requested(), 1);
659    }
660
661    // Check that the number of times requested equals to the number of headers divided by request
662    // limit.
663    #[tokio::test]
664    async fn requests_correct_number_of_times() {
665        // Generate some random blocks
666        let factory = create_test_provider_factory();
667        let mut rng = generators::rng();
668        let blocks = random_block_range(
669            &mut rng,
670            0..=199,
671            BlockRangeParams { parent: Some(B256::ZERO), tx_count: 1..2, ..Default::default() },
672        );
673
674        let headers = blocks.iter().map(|block| block.clone_sealed_header()).collect::<Vec<_>>();
675        let bodies = blocks
676            .into_iter()
677            .map(|block| (block.hash(), block.into_body()))
678            .collect::<HashMap<_, _>>();
679
680        insert_headers(&factory, &headers);
681
682        let request_limit = 10;
683        let client = Arc::new(TestBodiesClient::default().with_bodies(bodies.clone()));
684
685        let mut downloader = BodiesDownloaderBuilder::default()
686            .with_request_limit(request_limit)
687            .build::<reth_ethereum_primitives::Block, _, _>(
688            client.clone(),
689            Arc::new(TestConsensus::default()),
690            factory,
691        );
692        downloader.set_download_range(0..=199).expect("failed to set download range");
693
694        let _ = downloader.collect::<Vec<_>>().await;
695        assert_eq!(client.times_requested(), 20);
696    }
697
698    // Check that bodies are returned in correct order
699    // after resetting the download range multiple times.
700    #[tokio::test]
701    async fn streams_bodies_in_order_after_range_reset() {
702        // Generate some random blocks
703        let factory = create_test_provider_factory();
704        let (headers, mut bodies) = generate_bodies(0..=99);
705
706        insert_headers(&factory, &headers);
707
708        let stream_batch_size = 20;
709        let request_limit = 10;
710        let client = Arc::new(
711            TestBodiesClient::default().with_bodies(bodies.clone()).with_should_delay(true),
712        );
713        let mut downloader = BodiesDownloaderBuilder::default()
714            .with_stream_batch_size(stream_batch_size)
715            .with_request_limit(request_limit)
716            .build::<reth_ethereum_primitives::Block, _, _>(
717                client.clone(),
718                Arc::new(TestConsensus::default()),
719                factory,
720            );
721
722        let mut range_start = 0;
723        while range_start < 100 {
724            downloader.set_download_range(range_start..=99).expect("failed to set download range");
725
726            assert_matches!(
727                downloader.next().await,
728                Some(Ok(res)) => assert_eq!(res, zip_blocks(headers.iter().skip(range_start as usize).take(stream_batch_size), &mut bodies))
729            );
730            assert!(downloader.latest_queued_block_number >= Some(range_start));
731            range_start += stream_batch_size as u64;
732        }
733    }
734
735    // Check that the downloader picks up the new range and downloads bodies after previous range
736    // was completed.
737    #[tokio::test]
738    async fn can_download_new_range_after_termination() {
739        // Generate some random blocks
740        let factory = create_test_provider_factory();
741        let (headers, mut bodies) = generate_bodies(0..=199);
742
743        insert_headers(&factory, &headers);
744
745        let client = Arc::new(TestBodiesClient::default().with_bodies(bodies.clone()));
746
747        let mut downloader = BodiesDownloaderBuilder::default()
748            .with_stream_batch_size(100)
749            .build::<reth_ethereum_primitives::Block, _, _>(
750            client.clone(),
751            Arc::new(TestConsensus::default()),
752            factory,
753        );
754
755        // Set and download the first range
756        downloader.set_download_range(0..=99).expect("failed to set download range");
757        assert_matches!(
758            downloader.next().await,
759            Some(Ok(res)) => assert_eq!(res, zip_blocks(headers.iter().take(100), &mut bodies))
760        );
761
762        // Check that the stream is terminated
763        assert!(downloader.next().await.is_none());
764
765        // Set and download the second range
766        downloader.set_download_range(100..=199).expect("failed to set download range");
767        assert_matches!(
768            downloader.next().await,
769            Some(Ok(res)) => assert_eq!(res, zip_blocks(headers.iter().skip(100), &mut bodies))
770        );
771    }
772
773    // Check that the downloader continues after the size limit is reached.
774    #[tokio::test]
775    async fn can_download_after_exceeding_limit() {
776        // Generate some random blocks
777        let factory = create_test_provider_factory();
778        let (headers, mut bodies) = generate_bodies(0..=199);
779
780        insert_headers(&factory, &headers);
781
782        let client = Arc::new(TestBodiesClient::default().with_bodies(bodies.clone()));
783
784        // Set the max buffered block size to 1 byte, to make sure that every response exceeds the
785        // limit
786        let mut downloader = BodiesDownloaderBuilder::default()
787            .with_stream_batch_size(10)
788            .with_request_limit(1)
789            .with_max_buffered_blocks_size_bytes(1)
790            .build::<reth_ethereum_primitives::Block, _, _>(
791                client.clone(),
792                Arc::new(TestConsensus::default()),
793                factory,
794            );
795
796        // Set and download the entire range
797        downloader.set_download_range(0..=199).expect("failed to set download range");
798        let mut header = 0;
799        while let Some(Ok(resp)) = downloader.next().await {
800            assert_eq!(resp, zip_blocks(headers.iter().skip(header).take(resp.len()), &mut bodies));
801            header += resp.len();
802        }
803    }
804
805    // Check that the downloader can tolerate a few completely empty responses
806    #[tokio::test]
807    async fn can_tolerate_empty_responses() {
808        // Generate some random blocks
809        let factory = create_test_provider_factory();
810        let (headers, mut bodies) = generate_bodies(0..=99);
811
812        insert_headers(&factory, &headers);
813
814        // respond with empty bodies for every other request.
815        let client = Arc::new(
816            TestBodiesClient::default().with_bodies(bodies.clone()).with_empty_responses(2),
817        );
818
819        let mut downloader = BodiesDownloaderBuilder::default()
820            .with_request_limit(3)
821            .with_stream_batch_size(100)
822            .build::<reth_ethereum_primitives::Block, _, _>(
823                client.clone(),
824                Arc::new(TestConsensus::default()),
825                factory,
826            );
827
828        // Download the requested range
829        downloader.set_download_range(0..=99).expect("failed to set download range");
830        assert_matches!(
831            downloader.next().await,
832            Some(Ok(res)) => assert_eq!(res, zip_blocks(headers.iter().take(100), &mut bodies))
833        );
834    }
835}
reth_downloaders/bodies/bodies.rs

reth_downloaders/bodies/
bodies.rs
