Skip to main content

ab_networking/utils/
piece_provider.rs

1//! Provides methods to retrieve pieces from DSN.
2
3#[cfg(test)]
4mod tests;
5
6use crate::protocols::request_response::handlers::cached_piece_by_index::{
7    CachedPieceByIndexRequest, CachedPieceByIndexResponse, PieceResult,
8};
9use crate::protocols::request_response::handlers::piece_by_index::{
10    PieceByIndexRequest, PieceByIndexResponse,
11};
12use crate::utils::multihash::ToMultihash;
13use crate::{Multihash, Node};
14use ab_core_primitives::pieces::{Piece, PieceIndex};
15use async_lock::{Semaphore, SemaphoreGuard};
16use async_trait::async_trait;
17use futures::channel::mpsc;
18use futures::future::FusedFuture;
19use futures::stream::FuturesUnordered;
20use futures::{FutureExt, Stream, StreamExt, stream};
21use libp2p::kad::{K_VALUE, KBucketDistance, KBucketKey, RecordKey};
22use libp2p::{Multiaddr, PeerId};
23use rand::prelude::*;
24use std::any::type_name;
25use std::collections::{HashMap, HashSet};
26use std::fmt;
27use std::future::ready;
28use std::pin::Pin;
29use std::sync::Arc as StdArc;
30use std::task::Poll;
31use tokio_stream::StreamMap;
32use tracing::{Instrument, debug, trace, warn};
33use triomphe::Arc;
34
35/// Validates piece against using its root.
36#[async_trait]
37pub trait PieceValidator: Sync + Send {
38    /// Validates piece against using its root.
39    async fn validate_piece(
40        &self,
41        source_peer_id: PeerId,
42        piece_index: PieceIndex,
43        piece: Piece,
44    ) -> Option<Piece>;
45}
46
47/// Stub implementation for piece validation.
48#[derive(Debug, Clone, Copy)]
49pub struct NoPieceValidator;
50
51#[async_trait]
52impl PieceValidator for NoPieceValidator {
53    async fn validate_piece(&self, _: PeerId, _: PieceIndex, piece: Piece) -> Option<Piece> {
54        Some(piece)
55    }
56}
57
58/// Piece provider with cancellation and piece validator.
59/// Use `NoPieceValidator` to disable validation.
60#[derive(Clone)]
61pub struct PieceProvider<PV> {
62    node: Node,
63    piece_validator: PV,
64    piece_downloading_semaphore: Arc<Semaphore>,
65}
66
67impl<PV> fmt::Debug for PieceProvider<PV> {
68    #[inline]
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        f.debug_struct(&format!("PieceProvider<{}>", type_name::<PV>()))
71            .finish_non_exhaustive()
72    }
73}
74
75impl<PV> PieceProvider<PV>
76where
77    PV: PieceValidator,
78{
79    /// Creates new piece provider.
80    pub fn new(
81        node: Node,
82        piece_validator: PV,
83        piece_downloading_semaphore: Arc<Semaphore>,
84    ) -> Self {
85        Self {
86            node,
87            piece_validator,
88            piece_downloading_semaphore,
89        }
90    }
91
92    /// Get pieces with provided indices from cache.
93    ///
94    /// Number of elements in returned stream is the same as number of unique `piece_indices`.
95    pub fn get_from_cache<'a, PieceIndices>(
96        &'a self,
97        piece_indices: PieceIndices,
98    ) -> impl Future<Output = impl Stream<Item = (PieceIndex, Option<Piece>)> + Unpin + 'a>
99    where
100        PieceIndices: IntoIterator<Item = PieceIndex> + 'a,
101    {
102        let download_id = rand::random::<u64>();
103        let (tx, mut rx) = mpsc::unbounded();
104        let fut = async move {
105            let not_downloaded_pieces = download_cached_pieces(
106                piece_indices,
107                &self.node,
108                &self.piece_validator,
109                &tx,
110                &self.piece_downloading_semaphore,
111            )
112            .await;
113
114            if not_downloaded_pieces.is_empty() {
115                debug!("Done");
116                return;
117            }
118
119            for piece_index in not_downloaded_pieces {
120                tx.unbounded_send((piece_index, None))
121                    .expect("This future isn't polled after receiver is dropped; qed");
122            }
123
124            debug!("Done #2");
125        };
126
127        let mut fut = Box::pin(fut.instrument(tracing::info_span!("", %download_id)).fuse());
128
129        // Drive above future and stream back any pieces that were downloaded so far
130        ready(stream::poll_fn(move |cx| {
131            if !fut.is_terminated() {
132                // Result doesn't matter, we'll need to poll stream below anyway
133                let _: Poll<()> = fut.poll_unpin(cx);
134            }
135
136            if let Poll::Ready(maybe_result) = rx.poll_next_unpin(cx) {
137                return Poll::Ready(maybe_result);
138            }
139
140            // Exit will be done by the stream above
141            Poll::Pending
142        }))
143    }
144
145    /// Returns piece by its index from farmer's piece cache (L2)
146    pub async fn get_piece_from_cache(&self, piece_index: PieceIndex) -> Option<Piece> {
147        let key = RecordKey::from(piece_index.to_multihash());
148
149        let request_batch = self.node.get_requests_batch_handle().await;
150        let mut get_providers_stream = request_batch
151            .get_providers(key.clone())
152            .await
153            .inspect_err(|err| warn!(%piece_index,?key, ?err, "get_providers returned an error"))
154            .ok()?;
155
156        let key = hex::encode(&key);
157        while let Some(provider_id) = get_providers_stream.next().await {
158            trace!(%piece_index, key, %provider_id, "get_providers returned an item");
159
160            let Ok(PieceByIndexResponse {
161                piece,
162                cached_pieces: _,
163            }) = request_batch
164                .send_generic_request(
165                    provider_id,
166                    Vec::new(),
167                    PieceByIndexRequest {
168                        piece_index,
169                        cached_pieces: StdArc::default(),
170                    },
171                )
172                .await
173                .inspect_err(
174                    |error| debug!(%piece_index, key, %provider_id, ?error, "Piece request failed"),
175                )
176            else {
177                continue;
178            };
179
180            if let Some(piece) = piece {
181                trace!(%piece_index, key, %provider_id, "Piece request succeeded");
182
183                return self
184                    .piece_validator
185                    .validate_piece(provider_id, piece_index, piece)
186                    .await;
187            }
188
189            debug!(%piece_index, key, %provider_id, "Piece request returned empty piece");
190        }
191
192        None
193    }
194
195    /// Get piece from a particular peer.
196    pub async fn get_piece_from_peer(
197        &self,
198        peer_id: PeerId,
199        piece_index: PieceIndex,
200    ) -> Option<Piece> {
201        // TODO: Take advantage of `cached_pieces`
202        let PieceByIndexResponse {
203            piece,
204            cached_pieces: _,
205        } = self
206            .node
207            .send_generic_request(
208                peer_id,
209                Vec::new(),
210                PieceByIndexRequest {
211                    piece_index,
212                    cached_pieces: StdArc::default(),
213                },
214            )
215            .await
216            .inspect_err(|error| debug!(%peer_id, %piece_index, ?error, "Piece request failed"))
217            .ok()?;
218
219        if let Some(piece) = piece {
220            trace!(%peer_id, %piece_index, "Piece request succeeded");
221
222            return self
223                .piece_validator
224                .validate_piece(peer_id, piece_index, piece)
225                .await;
226        }
227
228        debug!(%peer_id, %piece_index, "Piece request returned empty piece");
229
230        None
231    }
232
233    /// Get piece from archival storage (L1). The algorithm tries to get a piece from currently
234    /// connected peers and falls back to random walking.
235    pub async fn get_piece_from_archival_storage(
236        &self,
237        piece_index: PieceIndex,
238        max_random_walking_rounds: usize,
239    ) -> Option<Piece> {
240        // TODO: consider using retry policy for L1 lookups as well.
241        trace!(%piece_index, "Getting piece from archival storage..");
242
243        let connected_servers = {
244            let connected_servers = match self.node.connected_servers().await {
245                Ok(connected_servers) => connected_servers,
246                Err(err) => {
247                    debug!(%piece_index, ?err, "Cannot get connected peers (DSN L1 lookup)");
248
249                    Vec::new()
250                }
251            };
252
253            HashSet::<PeerId>::from_iter(connected_servers)
254        };
255
256        if connected_servers.is_empty() {
257            debug!(%piece_index, "Cannot acquire piece from no connected peers (DSN L1 lookup)");
258        } else {
259            for peer_id in connected_servers {
260                let maybe_piece = self.get_piece_from_peer(peer_id, piece_index).await;
261
262                if maybe_piece.is_some() {
263                    trace!(%piece_index, %peer_id, "DSN L1 lookup from connected peers succeeded");
264
265                    return maybe_piece;
266                }
267            }
268        }
269
270        trace!(%piece_index, "Getting piece from DSN L1 using random walk.");
271        let random_walk_result = self
272            .get_piece_by_random_walking(piece_index, max_random_walking_rounds)
273            .await;
274
275        if random_walk_result.is_some() {
276            trace!(%piece_index, "DSN L1 lookup via random walk succeeded");
277
278            return random_walk_result;
279        }
280
281        debug!(
282            %piece_index,
283            %max_random_walking_rounds,
284            "Cannot acquire piece from DSN L1: random walk failed"
285        );
286
287        None
288    }
289
290    /// Get piece from L1 by random walking
291    async fn get_piece_by_random_walking(
292        &self,
293        piece_index: PieceIndex,
294        walking_rounds: usize,
295    ) -> Option<Piece> {
296        for round in 0..walking_rounds {
297            debug!(%piece_index, round, "Random walk round");
298
299            let result = self
300                .get_piece_by_random_walking_from_single_round(piece_index, round)
301                .await;
302
303            if result.is_some() {
304                return result;
305            }
306        }
307
308        debug!(%piece_index, "Random walking piece retrieval failed.");
309
310        None
311    }
312
313    /// Get piece from L1 by random walking (single round)
314    async fn get_piece_by_random_walking_from_single_round(
315        &self,
316        piece_index: PieceIndex,
317        round: usize,
318    ) -> Option<Piece> {
319        // TODO: Take advantage of `cached_pieces`
320        trace!(%piece_index, "get_piece_by_random_walking round");
321
322        // Random walk key
323        let key = PeerId::random();
324
325        let request_batch = self.node.get_requests_batch_handle().await;
326        let mut get_closest_peers_stream = request_batch
327            .get_closest_peers(key.into())
328            .await
329            .inspect_err(|err| warn!(%piece_index, ?key, ?err, %round, "get_closest_peers returned an error"))
330            .ok()?;
331
332        while let Some(peer_id) = get_closest_peers_stream.next().await {
333            trace!(%piece_index, %peer_id, %round, "get_closest_peers returned an item");
334
335            let Ok(PieceByIndexResponse {
336                piece,
337                cached_pieces: _,
338            }) = request_batch
339                .send_generic_request(
340                    peer_id,
341                    Vec::new(),
342                    PieceByIndexRequest {
343                        piece_index,
344                        cached_pieces: StdArc::default(),
345                    },
346                )
347                .await
348                .inspect_err(
349                    |error| debug!(%peer_id, %piece_index, ?key, %round, ?error, "Piece request failed."),
350                )
351            else {
352                continue;
353            };
354
355            if let Some(piece) = piece {
356                trace!(%peer_id, %piece_index, ?key, %round,  "Piece request succeeded.");
357
358                return self
359                    .piece_validator
360                    .validate_piece(peer_id, piece_index, piece)
361                    .await;
362            }
363
364            debug!(%peer_id, %piece_index, ?key, %round, "Piece request returned empty piece.");
365        }
366
367        None
368    }
369}
370
371/// Up to [`K_VALUE`] closest known peers of a piece with their addresses
372struct KnownClosestPeers {
373    key: KBucketKey<Multihash>,
374    /// Sorted in ascending distance order, contains at most [`K_VALUE`] peers
375    peers: Vec<(KBucketDistance, PeerId, Vec<Multiaddr>)>,
376}
377
378impl KnownClosestPeers {
379    fn new(key: KBucketKey<Multihash>) -> Self {
380        Self {
381            key,
382            peers: Vec::with_capacity(K_VALUE.get()),
383        }
384    }
385
386    fn add_peer(&mut self, local_peer_id: &PeerId, peer_id: PeerId, addresses: Vec<Multiaddr>) {
387        // Kademlia doesn't store the local peer or peers without addresses either
388        if peer_id == *local_peer_id || addresses.is_empty() {
389            return;
390        }
391
392        if let Some((_distance, _peer_id, known_addresses)) = self
393            .peers
394            .iter_mut()
395            .find(|(_distance, known_peer_id, _addresses)| *known_peer_id == peer_id)
396        {
397            for address in addresses {
398                if !known_addresses.contains(&address) {
399                    known_addresses.push(address);
400                }
401            }
402            return;
403        }
404
405        let distance = KBucketKey::from(peer_id).distance(&self.key);
406        let position = self
407            .peers
408            .partition_point(|(known_distance, _peer_id, _addresses)| *known_distance < distance);
409        if position >= K_VALUE.get() {
410            return;
411        }
412        if self.peers.len() == K_VALUE.get() {
413            self.peers.pop();
414        }
415        self.peers.insert(position, (distance, peer_id, addresses));
416    }
417
418    /// Returned peers are sorted in ascending distance order
419    fn iter(&self) -> impl Iterator<Item = (&PeerId, &Vec<Multiaddr>)> {
420        self.peers
421            .iter()
422            .map(|(_distance, peer_id, addresses)| (peer_id, addresses))
423    }
424}
425
426/// Takes pieces to download as an input, sends results with pieces that were downloaded
427/// successfully and returns those that were not downloaded
428async fn download_cached_pieces<PV, PieceIndices>(
429    piece_indices: PieceIndices,
430    node: &Node,
431    piece_validator: &PV,
432    results: &mpsc::UnboundedSender<(PieceIndex, Option<Piece>)>,
433    semaphore: &Semaphore,
434) -> impl ExactSizeIterator<Item = PieceIndex>
435where
436    PV: PieceValidator,
437    PieceIndices: IntoIterator<Item = PieceIndex>,
438{
439    // Make sure every piece index has an entry since this will be the primary container for
440    // tracking pieces to download going forward.
441    //
442    // At the end pieces that were not downloaded will remain with a collection of known closest
443    // peers for them.
444    let mut pieces_to_download = piece_indices
445        .into_iter()
446        .map(|piece_index| async move {
447            let key = piece_index.to_multihash();
448            let mut closest_peers = KnownClosestPeers::new(KBucketKey::from(key));
449
450            let local_closest_peers = node
451                .get_closest_local_peers(key, None)
452                .await
453                .unwrap_or_default();
454
455            // Seed with local closest peers
456            for (peer_id, addresses) in local_closest_peers {
457                closest_peers.add_peer(&node.id(), peer_id, addresses);
458            }
459
460            (piece_index, closest_peers)
461        })
462        .collect::<FuturesUnordered<_>>()
463        .collect::<HashMap<_, _>>()
464        .await;
465
466    let num_pieces = pieces_to_download.len();
467    debug!(%num_pieces, "Starting");
468
469    let mut checked_peers = HashSet::new();
470
471    let Ok(connected_servers) = node.connected_servers().await else {
472        trace!("Connected servers error");
473        return pieces_to_download.into_keys();
474    };
475
476    let num_connected_servers = connected_servers.len();
477    debug!(
478        %num_connected_servers,
479        %num_pieces,
480        "Starting downloading"
481    );
482
483    // Dispatch initial set of requests to peers with checked pieces distributed uniformly
484    let mut downloading_stream = connected_servers
485        .into_iter()
486        .take(num_pieces)
487        .enumerate()
488        .map(|(peer_index, peer_id)| {
489            checked_peers.insert(peer_id);
490
491            // Inside to avoid division by zero in case there are no connected servers or pieces
492            let step = num_pieces / num_connected_servers.min(num_pieces);
493
494            // Take unique first piece index for each connected peer and the rest just to check
495            // cached pieces up to recommended limit
496            let mut check_cached_pieces = pieces_to_download
497                .keys()
498                .cycle()
499                .skip(step * peer_index)
500                // + 1 because one index below is removed below
501                .take(num_pieces.min(CachedPieceByIndexRequest::RECOMMENDED_LIMIT + 1))
502                .copied()
503                .collect::<Vec<_>>();
504            // Pick first piece index as the piece we want to download
505            let piece_index = check_cached_pieces.swap_remove(0);
506
507            trace!(%peer_id, %piece_index, "Downloading piece from initially connected peer");
508
509            let permit = semaphore.try_acquire();
510
511            let fut = async move {
512                let permit = match permit {
513                    Some(permit) => permit,
514                    None => semaphore.acquire().await,
515                };
516
517                download_cached_piece_from_peer(
518                    node,
519                    piece_validator,
520                    peer_id,
521                    Vec::new(),
522                    StdArc::new(check_cached_pieces),
523                    piece_index,
524                    HashSet::new(),
525                    HashSet::new(),
526                    permit,
527                )
528                .await
529            };
530
531            (piece_index, Box::pin(fut.into_stream()) as _)
532        })
533        .collect::<StreamMap<_, _>>();
534
535    loop {
536        // Process up to 50% of the pieces concurrently
537        let mut additional_pieces_to_download =
538            (num_pieces / 2).saturating_sub(downloading_stream.len());
539        if additional_pieces_to_download > 0 {
540            trace!(
541                %additional_pieces_to_download,
542                num_pieces,
543                currently_downloading = %downloading_stream.len(),
544                "Downloading additional pieces from closest peers"
545            );
546            // Pick up any newly connected peers (if any)
547            'outer: for peer_id in node
548                .connected_servers()
549                .await
550                .unwrap_or_default()
551                .into_iter()
552                .filter(|peer_id| checked_peers.insert(*peer_id))
553                .take(additional_pieces_to_download)
554            {
555                let permit = if downloading_stream.is_empty() {
556                    semaphore.acquire().await
557                } else if let Some(permit) = semaphore.try_acquire() {
558                    permit
559                } else {
560                    break;
561                };
562
563                for &piece_index in pieces_to_download.keys() {
564                    if downloading_stream.contains_key(&piece_index) {
565                        continue;
566                    }
567
568                    trace!(%peer_id, %piece_index, "Downloading piece from newly connected peer");
569
570                    let check_cached_pieces = sample_cached_piece_indices(
571                        pieces_to_download.keys(),
572                        &HashSet::new(),
573                        &HashSet::new(),
574                        piece_index,
575                    );
576                    let fut = download_cached_piece_from_peer(
577                        node,
578                        piece_validator,
579                        peer_id,
580                        Vec::new(),
581                        StdArc::new(check_cached_pieces),
582                        piece_index,
583                        HashSet::new(),
584                        HashSet::new(),
585                        permit,
586                    );
587
588                    downloading_stream.insert(piece_index, Box::pin(fut.into_stream()) as _);
589                    additional_pieces_to_download -= 1;
590
591                    continue 'outer;
592                }
593
594                break;
595            }
596
597            // Pick up more pieces to download from the closest peers
598            // Ideally we'd not allocate here, but it is hard to explain to the compiler that
599            // entries are not removed otherwise
600            let pieces_indices_to_download = pieces_to_download.keys().copied().collect::<Vec<_>>();
601            for piece_index in pieces_indices_to_download {
602                if additional_pieces_to_download == 0 {
603                    break;
604                }
605                if downloading_stream.contains_key(&piece_index) {
606                    continue;
607                }
608                let permit = if downloading_stream.is_empty() {
609                    semaphore.acquire().await
610                } else if let Some(permit) = semaphore.try_acquire() {
611                    permit
612                } else {
613                    break;
614                };
615
616                let closest_peers_to_check = pieces_to_download
617                    .get(&piece_index)
618                    .expect("Entries are not removed here; qed");
619                for (&peer_id, addresses) in closest_peers_to_check.iter() {
620                    if !checked_peers.insert(peer_id) {
621                        continue;
622                    }
623
624                    trace!(%peer_id, %piece_index, "Downloading piece from closest peer");
625
626                    let check_cached_pieces = sample_cached_piece_indices(
627                        pieces_to_download.keys(),
628                        &HashSet::new(),
629                        &HashSet::new(),
630                        piece_index,
631                    );
632                    let fut = download_cached_piece_from_peer(
633                        node,
634                        piece_validator,
635                        peer_id,
636                        addresses.clone(),
637                        StdArc::new(check_cached_pieces),
638                        piece_index,
639                        HashSet::new(),
640                        HashSet::new(),
641                        permit,
642                    );
643
644                    downloading_stream.insert(piece_index, Box::pin(fut.into_stream()) as _);
645                    additional_pieces_to_download -= 1;
646                    break;
647                }
648            }
649
650            trace!(
651                pieces_left = %additional_pieces_to_download,
652                "Initiated downloading additional pieces from closest peers"
653            );
654        }
655
656        let Some((piece_index, result)) = downloading_stream.next().await else {
657            if !pieces_to_download.is_empty() {
658                debug!(
659                    %num_pieces,
660                    to_download = %pieces_to_download.len(),
661                    "Finished downloading early"
662                );
663                // Nothing was downloaded, we're done here
664                break;
665            }
666            break;
667        };
668        process_downloading_result(
669            piece_index,
670            result,
671            &mut pieces_to_download,
672            &mut downloading_stream,
673            node,
674            piece_validator,
675            results,
676        );
677
678        if pieces_to_download.is_empty() {
679            break;
680        }
681    }
682
683    pieces_to_download.into_keys()
684}
685
686fn process_downloading_result<'a, 'b, PV>(
687    piece_index: PieceIndex,
688    result: DownloadedPieceFromPeer<'a>,
689    pieces_to_download: &'b mut HashMap<PieceIndex, KnownClosestPeers>,
690    downloading_stream: &'b mut StreamMap<
691        PieceIndex,
692        Pin<Box<dyn Stream<Item = DownloadedPieceFromPeer<'a>> + Send + 'a>>,
693    >,
694    node: &'a Node,
695    piece_validator: &'a PV,
696    results: &'a mpsc::UnboundedSender<(PieceIndex, Option<Piece>)>,
697) where
698    PV: PieceValidator,
699{
700    let DownloadedPieceFromPeer {
701        peer_id,
702        result,
703        mut cached_pieces,
704        not_cached_pieces,
705        permit,
706    } = result;
707    trace!(%piece_index, %peer_id, result = %result.is_some(), "Piece response");
708
709    let Some(result) = result else {
710        // Downloading failed, ignore peer
711        return;
712    };
713
714    match result {
715        PieceResult::Piece(piece) => {
716            trace!(%piece_index, %peer_id, "Got piece");
717
718            // Downloaded successfully
719            pieces_to_download.remove(&piece_index);
720
721            results
722                .unbounded_send((piece_index, Some(piece)))
723                .expect("This future isn't polled after receiver is dropped; qed");
724
725            if pieces_to_download.is_empty() {
726                return;
727            }
728
729            cached_pieces.remove(&piece_index);
730        }
731        PieceResult::ClosestPeers(closest_peers) => {
732            trace!(%piece_index, %peer_id, "Got closest peers");
733
734            // Store closer peers in case piece index was not downloaded yet
735            if let Some(known_closest_peers) = pieces_to_download.get_mut(&piece_index) {
736                for (peer_id, addresses) in Vec::from(closest_peers) {
737                    known_closest_peers.add_peer(&node.id(), peer_id, addresses);
738                }
739            }
740
741            // No need to ask this peer again if they claimed to have this piece index earlier
742            if cached_pieces.remove(&piece_index) {
743                return;
744            }
745        }
746    }
747
748    let mut maybe_piece_index_to_download_next = None;
749    // Clear useless entries in cached pieces and find something to download next
750    cached_pieces.retain(|piece_index| {
751        // Clear downloaded pieces
752        if !pieces_to_download.contains_key(piece_index) {
753            return false;
754        }
755
756        // Try to pick a piece to download that is not being downloaded already
757        if maybe_piece_index_to_download_next.is_none()
758            && !downloading_stream.contains_key(piece_index)
759        {
760            maybe_piece_index_to_download_next.replace(*piece_index);
761            // We'll check it later when receiving response
762            return true;
763        }
764
765        // Retain everything else
766        true
767    });
768
769    let piece_index_to_download_next = if let Some(piece_index) = maybe_piece_index_to_download_next
770    {
771        trace!(%piece_index, %peer_id, "Next piece to download from peer");
772        piece_index
773    } else {
774        trace!(%peer_id, "Peer doesn't have anything else");
775        // Nothing left to do with this peer
776        return;
777    };
778
779    let fut = download_cached_piece_from_peer(
780        node,
781        piece_validator,
782        peer_id,
783        Vec::new(),
784        // Sample more random cached piece indices for connected peer, algorithm can be
785        // improved, but has to be something simple and this should do it for now
786        StdArc::new(sample_cached_piece_indices(
787            pieces_to_download.keys(),
788            &cached_pieces,
789            &not_cached_pieces,
790            piece_index_to_download_next,
791        )),
792        piece_index_to_download_next,
793        cached_pieces,
794        not_cached_pieces,
795        permit,
796    );
797    downloading_stream.insert(piece_index_to_download_next, Box::pin(fut.into_stream()));
798}
799
800fn sample_cached_piece_indices<'a, I>(
801    pieces_to_download: I,
802    cached_pieces: &HashSet<PieceIndex>,
803    not_cached_pieces: &HashSet<PieceIndex>,
804    piece_index_to_download_next: PieceIndex,
805) -> Vec<PieceIndex>
806where
807    I: Iterator<Item = &'a PieceIndex>,
808{
809    pieces_to_download
810        // Do a bit of work to filter-out piece indices we already know remote peer
811        // has or doesn't to decrease burden on them
812        .filter_map(|piece_index| {
813            if piece_index == &piece_index_to_download_next
814                || cached_pieces.contains(piece_index)
815                || not_cached_pieces.contains(piece_index)
816            {
817                None
818            } else {
819                Some(*piece_index)
820            }
821        })
822        .sample(
823            &mut rand::rng(),
824            CachedPieceByIndexRequest::RECOMMENDED_LIMIT,
825        )
826}
827
828struct DownloadedPieceFromPeer<'a> {
829    peer_id: PeerId,
830    result: Option<PieceResult>,
831    cached_pieces: HashSet<PieceIndex>,
832    not_cached_pieces: HashSet<PieceIndex>,
833    permit: SemaphoreGuard<'a>,
834}
835
836/// `check_cached_pieces` contains a list of pieces for peer to filter-out according to locally
837/// caches pieces, `cached_pieces` and `not_cached_pieces` contain piece indices peer claims is
838/// known to have or not have already
839#[expect(clippy::too_many_arguments, reason = "Internal API")]
840async fn download_cached_piece_from_peer<'a, PV>(
841    node: &'a Node,
842    piece_validator: &'a PV,
843    peer_id: PeerId,
844    addresses: Vec<Multiaddr>,
845    check_cached_pieces: StdArc<Vec<PieceIndex>>,
846    piece_index: PieceIndex,
847    mut cached_pieces: HashSet<PieceIndex>,
848    mut not_cached_pieces: HashSet<PieceIndex>,
849    permit: SemaphoreGuard<'a>,
850) -> DownloadedPieceFromPeer<'a>
851where
852    PV: PieceValidator,
853{
854    let result = match node
855        .send_generic_request(
856            peer_id,
857            addresses,
858            CachedPieceByIndexRequest {
859                piece_index,
860                cached_pieces: StdArc::clone(&check_cached_pieces),
861            },
862        )
863        .await
864    {
865        Ok(response) => {
866            let CachedPieceByIndexResponse {
867                result,
868                cached_pieces,
869            } = response;
870
871            match result {
872                PieceResult::Piece(piece) => piece_validator
873                    .validate_piece(peer_id, piece_index, piece)
874                    .await
875                    .map(|piece| CachedPieceByIndexResponse {
876                        result: PieceResult::Piece(piece),
877                        cached_pieces,
878                    }),
879                PieceResult::ClosestPeers(closest_peers) => Some(CachedPieceByIndexResponse {
880                    result: PieceResult::ClosestPeers(closest_peers),
881                    cached_pieces,
882                }),
883            }
884        }
885        Err(error) => {
886            debug!(%error, %peer_id, %piece_index, "Failed to download cached piece from peer");
887
888            None
889        }
890    };
891
892    match result {
893        Some(result) => {
894            cached_pieces.extend(result.cached_pieces);
895            not_cached_pieces.extend(
896                check_cached_pieces
897                    .iter()
898                    .filter(|piece_index| !cached_pieces.contains(piece_index))
899                    .copied(),
900            );
901
902            DownloadedPieceFromPeer {
903                peer_id,
904                result: Some(result.result),
905                cached_pieces: { cached_pieces },
906                not_cached_pieces,
907                permit,
908            }
909        }
910        None => DownloadedPieceFromPeer {
911            peer_id,
912            result: None,
913            cached_pieces,
914            not_cached_pieces,
915            permit,
916        },
917    }
918}