Skip to main content

ab_farmer/
farmer_cache.rs

1//! A container that caches pieces
2//!
3//! Farmer cache is a container that orchestrates a bunch of piece and plot caches that together
4//! persist pieces in a way that is easy to retrieve comparing to decoding pieces from plots.
5
6mod metrics;
7mod piece_cache_state;
8// TODO: Not supported under Miri: https://github.com/rust-lang/miri/issues/5137
9#[cfg(not(miri))]
10#[cfg(test)]
11mod tests;
12
13use crate::farm::{MaybePieceStoredResult, PieceCache, PieceCacheId, PieceCacheOffset, PlotCache};
14use crate::farmer_cache::metrics::FarmerCacheMetrics;
15use crate::farmer_cache::piece_cache_state::PieceCachesState;
16use crate::node_client::NodeClient;
17use crate::utils::run_future_in_dedicated_thread;
18use ab_core_primitives::pieces::{Piece, PieceIndex};
19use ab_core_primitives::segments::SegmentIndex;
20use ab_data_retrieval::piece_getter::PieceGetter;
21use ab_networking::KeyWithDistance;
22use ab_networking::libp2p::PeerId;
23use ab_networking::libp2p::kad::RecordKey;
24use ab_networking::utils::multihash::ToMultihash;
25use async_lock::RwLock as AsyncRwLock;
26use bytesize::ByteSize;
27use event_listener_primitives::{Bag, HandlerId};
28use futures::channel::mpsc;
29use futures::future::{Either, FusedFuture};
30use futures::stream::{FuturesOrdered, FuturesUnordered};
31use futures::{FutureExt, SinkExt, Stream, StreamExt, select, stream};
32use parking_lot::RwLock;
33use prometheus_client::registry::Registry;
34use rand::prelude::*;
35use rayon::prelude::*;
36use std::collections::hash_map::Entry;
37use std::collections::{HashMap, HashSet};
38use std::future::join;
39use std::sync::Arc as StdArc;
40use std::sync::atomic::{AtomicUsize, Ordering};
41use std::task::Poll;
42use std::time::Duration;
43use std::{fmt, mem};
44use tokio::sync::Semaphore;
45use tokio::task::yield_now;
46use tracing::{Instrument, debug, error, info, info_span, trace, warn};
47use triomphe::Arc;
48
49const WORKER_CHANNEL_CAPACITY: usize = 100;
50const SYNC_BATCH_SIZE: usize = 256;
51const SYNC_CONCURRENT_BATCHES: usize = 4;
52/// Interval in pieces after which piece cache sync progress is logged
53const SYNC_PROGRESS_LOG_INTERVAL: usize = 100;
54const INITIAL_SYNC_FARM_INFO_CHECK_INTERVAL: Duration = Duration::from_secs(1);
55
56type HandlerFn<A> = StdArc<dyn Fn(&A) + Send + Sync + 'static>;
57type Handler<A> = Bag<HandlerFn<A>, A>;
58type CacheIndex = u8;
59
60#[derive(Default, Debug)]
61struct Handlers {
62    progress: Handler<f32>,
63}
64
65#[derive(Debug, Clone, Copy)]
66struct FarmerCacheOffset {
67    cache_index: CacheIndex,
68    piece_offset: PieceCacheOffset,
69}
70
71impl FarmerCacheOffset {
72    fn new(cache_index: CacheIndex, piece_offset: PieceCacheOffset) -> Self {
73        Self {
74            cache_index,
75            piece_offset,
76        }
77    }
78}
79
80#[derive(Debug, Clone)]
81struct CacheBackend {
82    backend: StdArc<dyn PieceCache>,
83    used_capacity: u32,
84    total_capacity: u32,
85}
86
87impl std::ops::Deref for CacheBackend {
88    type Target = StdArc<dyn PieceCache>;
89
90    fn deref(&self) -> &Self::Target {
91        &self.backend
92    }
93}
94
95impl CacheBackend {
96    fn new(backend: StdArc<dyn PieceCache>, total_capacity: u32) -> Self {
97        Self {
98            backend,
99            used_capacity: 0,
100            total_capacity,
101        }
102    }
103
104    fn next_free(&mut self) -> Option<PieceCacheOffset> {
105        let offset = self.used_capacity;
106        if offset < self.total_capacity {
107            self.used_capacity += 1;
108            Some(PieceCacheOffset(offset))
109        } else {
110            debug!(?offset, total_capacity = ?self.total_capacity, "No free space in cache backend");
111            None
112        }
113    }
114
115    fn free_size(&self) -> u32 {
116        self.total_capacity - self.used_capacity
117    }
118}
119
120#[derive(Debug)]
121struct CacheState {
122    cache_stored_pieces: HashMap<KeyWithDistance, FarmerCacheOffset>,
123    cache_free_offsets: Vec<FarmerCacheOffset>,
124    backend: CacheBackend,
125}
126
127#[derive(Debug)]
128enum WorkerCommand {
129    ReplaceBackingCaches {
130        new_piece_caches: Vec<StdArc<dyn PieceCache>>,
131    },
132    ForgetKey {
133        key: RecordKey,
134    },
135}
136
137/// Farmer cache worker used to drive the farmer cache backend
138#[derive(Debug)]
139#[must_use = "Farmer cache will not work unless its worker is running"]
140pub struct FarmerCacheWorker<NC>
141where
142    NC: fmt::Debug,
143{
144    peer_id: PeerId,
145    node_client: NC,
146    piece_caches: Arc<AsyncRwLock<PieceCachesState>>,
147    plot_caches: Arc<PlotCaches>,
148    handlers: Arc<Handlers>,
149    worker_receiver: Option<mpsc::Receiver<WorkerCommand>>,
150    metrics: Option<Arc<FarmerCacheMetrics>>,
151}
152
153impl<NC> FarmerCacheWorker<NC>
154where
155    NC: NodeClient,
156{
157    /// Run the cache worker with provided piece getter.
158    ///
159    /// NOTE: Piece getter must not depend on farmer cache in order to avoid reference cycles!
160    pub async fn run<PG>(mut self, piece_getter: PG)
161    where
162        PG: PieceGetter,
163    {
164        // Limit is dynamically set later
165        let mut last_segment_index_internal = SegmentIndex::ZERO;
166
167        let mut worker_receiver = self
168            .worker_receiver
169            .take()
170            .expect("Always set during worker instantiation");
171
172        if let Some(WorkerCommand::ReplaceBackingCaches { new_piece_caches }) =
173            worker_receiver.next().await
174        {
175            self.initialize(
176                &piece_getter,
177                &mut last_segment_index_internal,
178                new_piece_caches,
179            )
180            .await;
181        } else {
182            // Piece cache is dropped before backing caches were sent
183            return;
184        }
185
186        let mut super_segment_headers_notifications = match self
187            .node_client
188            .subscribe_new_super_segment_headers()
189            .await
190        {
191            Ok(super_segment_headers_notifications) => super_segment_headers_notifications,
192            Err(error) => {
193                error!(%error, "Failed to subscribe to new super segment headers notifications");
194                return;
195            }
196        };
197
198        // Keep up with segment indices that were potentially created since reinitialization
199        self.keep_up_after_initial_sync(&piece_getter, &mut last_segment_index_internal)
200            .await;
201
202        loop {
203            select! {
204                maybe_command = worker_receiver.next() => {
205                    let Some(command) = maybe_command else {
206                        // Nothing else left to do
207                        return;
208                    };
209
210                    self.handle_command(command, &piece_getter, &mut last_segment_index_internal).await;
211                }
212                maybe_super_segment_header = super_segment_headers_notifications.next().fuse() => {
213                    if let Some(super_segment_header) = maybe_super_segment_header {
214                        self
215                            .process_new_segment_index(
216                                &piece_getter,
217                                super_segment_header.max_segment_index.as_inner(),
218                                &mut last_segment_index_internal
219                            )
220                            .await;
221                    } else {
222                        // Keep-up sync only ends with subscription, which lasts for duration of an
223                        // instance
224                        return;
225                    }
226                }
227            }
228        }
229    }
230
231    async fn handle_command<PG>(
232        &self,
233        command: WorkerCommand,
234        piece_getter: &PG,
235        last_segment_index_internal: &mut SegmentIndex,
236    ) where
237        PG: PieceGetter,
238    {
239        match command {
240            WorkerCommand::ReplaceBackingCaches { new_piece_caches } => {
241                self.initialize(piece_getter, last_segment_index_internal, new_piece_caches)
242                    .await;
243            }
244            // TODO: Consider implementing optional re-sync of the piece instead of just forgetting
245            WorkerCommand::ForgetKey { key } => {
246                let mut caches = self.piece_caches.write().await;
247                let key = KeyWithDistance::new_with_record_key(self.peer_id, key);
248                let Some(offset) = caches.remove_stored_piece(&key) else {
249                    // Key not exist
250                    return;
251                };
252
253                let cache_index = offset.cache_index;
254                let piece_offset = offset.piece_offset;
255                let Some(backend) = caches.get_backend(cache_index).cloned() else {
256                    // Cache backend not exist
257                    return;
258                };
259
260                caches.push_dangling_free_offset(offset);
261                match backend.read_piece_index(piece_offset).await {
262                    Ok(Some(piece_index)) => {
263                        trace!(%piece_index, %cache_index, %piece_offset, "Forget piece");
264                    }
265                    Ok(None) => {
266                        warn!(
267                            %cache_index,
268                            %piece_offset,
269                            "Piece index out of range, this is likely an implementation bug, \
270                            not freeing heap element"
271                        );
272                    }
273                    Err(error) => {
274                        error!(
275                            %error,
276                            %cache_index,
277                            ?key,
278                            %piece_offset,
279                            "Error while reading piece from cache"
280                        );
281                    }
282                }
283            }
284        }
285    }
286
287    async fn initialize<PG>(
288        &self,
289        piece_getter: &PG,
290        last_segment_index_internal: &mut SegmentIndex,
291        new_piece_caches: Vec<StdArc<dyn PieceCache>>,
292    ) where
293        PG: PieceGetter,
294    {
295        info!("Initializing piece cache");
296
297        // Pull old cache state since it will be replaced with a new one and reuse its allocations
298        let (mut stored_pieces, mut dangling_free_offsets) =
299            mem::take(&mut *self.piece_caches.write().await).reuse();
300
301        debug!("Collecting pieces that were in the cache before");
302
303        if let Some(metrics) = &self.metrics {
304            metrics.piece_cache_capacity_total.set(0);
305            metrics.piece_cache_capacity_used.set(0);
306        }
307
308        let peer_id = self.peer_id;
309
310        // Build cache state of all backends
311        let piece_caches_number = new_piece_caches.len();
312        let maybe_caches_futures = new_piece_caches
313            .into_iter()
314            .enumerate()
315            .filter_map(|(cache_index, new_cache)| {
316                let total_capacity = new_cache.max_num_elements();
317                let mut backend = CacheBackend::new(new_cache, total_capacity);
318                let Ok(cache_index) = CacheIndex::try_from(cache_index) else {
319                    warn!(
320                        ?piece_caches_number,
321                        "Too many piece caches provided, {cache_index} cache will be ignored",
322                    );
323                    return None;
324                };
325
326                if let Some(metrics) = &self.metrics {
327                    metrics
328                        .piece_cache_capacity_total
329                        .inc_by(i64::from(total_capacity));
330                }
331
332                let init_fut = async move {
333                    let used_capacity = &mut backend.used_capacity;
334
335                    // Hack with first collecting into `Option` with `Option::take()` call
336                    // later is to satisfy compiler that gets confused about ownership
337                    // otherwise
338                    let mut maybe_contents = match backend.backend.contents().await {
339                        Ok(contents) => Some(contents),
340                        Err(error) => {
341                            warn!(%error, "Failed to get cache contents");
342
343                            None
344                        }
345                    };
346
347                    let mut cache_stored_pieces = HashMap::new();
348                    let mut cache_free_offsets = Vec::new();
349
350                    let Some(mut contents) = maybe_contents.take() else {
351                        drop(maybe_contents);
352
353                        return CacheState {
354                            cache_stored_pieces,
355                            cache_free_offsets,
356                            backend,
357                        };
358                    };
359
360                    while let Some(maybe_element_details) = contents.next().await {
361                        let (piece_offset, maybe_piece_index) = match maybe_element_details {
362                            Ok(element_details) => element_details,
363                            Err(error) => {
364                                warn!(%error, "Failed to get cache contents element details");
365                                break;
366                            }
367                        };
368                        let offset = FarmerCacheOffset::new(cache_index, piece_offset);
369                        match maybe_piece_index {
370                            Some(piece_index) => {
371                                *used_capacity = piece_offset.0 + 1;
372                                let record_key = RecordKey::from(piece_index.to_multihash());
373                                let key = KeyWithDistance::new_with_record_key(peer_id, record_key);
374                                cache_stored_pieces.insert(key, offset);
375                            }
376                            None => {
377                                // TODO: Optimize to not store all free offsets, only dangling
378                                //  offsets are actually necessary
379                                cache_free_offsets.push(offset);
380                            }
381                        }
382
383                        // Allow for task to be aborted
384                        yield_now().await;
385                    }
386
387                    drop(maybe_contents);
388                    drop(contents);
389
390                    CacheState {
391                        cache_stored_pieces,
392                        cache_free_offsets,
393                        backend,
394                    }
395                };
396
397                Some(run_future_in_dedicated_thread(
398                    move || init_fut.instrument(info_span!("", %cache_index)),
399                    format!("piece-cache.{cache_index}"),
400                ))
401            })
402            .collect::<Result<Vec<_>, _>>();
403
404        let caches_futures = match maybe_caches_futures {
405            Ok(caches_futures) => caches_futures,
406            Err(error) => {
407                error!(%error, "Failed to spawn piece cache reading thread");
408
409                return;
410            }
411        };
412
413        let mut backends = Vec::with_capacity(caches_futures.len());
414        let mut caches_futures = caches_futures.into_iter().collect::<FuturesOrdered<_>>();
415
416        while let Some(maybe_cache) = caches_futures.next().await {
417            match maybe_cache {
418                Ok(cache) => {
419                    let backend = cache.backend;
420                    for (key, cache_offset) in cache.cache_stored_pieces {
421                        if let Some(old_cache_offset) = stored_pieces.insert(key, cache_offset) {
422                            dangling_free_offsets.push_front(old_cache_offset);
423                        }
424                    }
425                    dangling_free_offsets.extend(
426                        cache.cache_free_offsets.into_iter().filter(|free_offset| {
427                            free_offset.piece_offset.0 < backend.used_capacity
428                        }),
429                    );
430                    backends.push(backend);
431                }
432                Err(_cancelled) => {
433                    error!("Piece cache reading thread panicked");
434
435                    return;
436                }
437            }
438        }
439
440        let mut caches = PieceCachesState::new(stored_pieces, dangling_free_offsets, backends);
441
442        info!("Synchronizing piece cache");
443
444        let last_segment_index = loop {
445            match self.node_client.farmer_app_info().await {
446                Ok(farmer_app_info) => {
447                    let last_segment_index =
448                        farmer_app_info.protocol_info.history_size.segment_index();
449                    // Wait for node to be either fully synced or to be aware of non-zero segment
450                    // index, which would indicate it has started DSN sync and knows about
451                    // up-to-date archived history.
452                    //
453                    // While this doesn't account for situations where node was offline for a long
454                    // time and is aware of old segment headers, this is good enough for piece cache
455                    // sync to proceed and should result in better user experience on average.
456                    if !farmer_app_info.syncing || last_segment_index > SegmentIndex::ZERO {
457                        break last_segment_index;
458                    }
459                }
460                Err(error) => {
461                    error!(
462                        %error,
463                        "Failed to get farmer app info from node, keeping old cache state without \
464                        updates"
465                    );
466
467                    // Not the latest, but at least something
468                    *self.piece_caches.write().await = caches;
469                    return;
470                }
471            }
472
473            tokio::time::sleep(INITIAL_SYNC_FARM_INFO_CHECK_INTERVAL).await;
474        };
475
476        debug!(%last_segment_index, "Identified last segment index");
477
478        // Collect all the piece indices that need to be stored, we will sort them later
479        let segment_indices = Vec::from_iter(SegmentIndex::ZERO..=last_segment_index);
480        // TODO: This may eventually be too much to store in memory, though right now it is a
481        //  non-issue in practice
482        let mut piece_indices_to_store = segment_indices
483            .into_par_iter()
484            .flat_map(|segment_index| {
485                segment_index
486                    .segment_piece_indexes()
487                    .into_par_iter()
488                    .map(|piece_index| {
489                        (
490                            KeyWithDistance::new(self.peer_id, piece_index.to_multihash()),
491                            piece_index,
492                        )
493                    })
494            })
495            .collect::<Vec<_>>();
496
497        // Sort pieces by distance from peer to piece such that they are in ascending order
498        // and have higher chance of download
499        piece_indices_to_store.par_sort_unstable_by(|(a_key, _), (b_key, _)| a_key.cmp(b_key));
500
501        // `HashMap` is faster than `BTreeMap`
502        let mut piece_indices_to_store = piece_indices_to_store
503            .into_iter()
504            .take(caches.total_capacity())
505            .collect::<HashMap<_, _>>();
506
507        let mut piece_caches_capacity_used = vec![0u32; caches.backends().len()];
508        // Filter-out piece indices that are stored, but should not be as well as clean
509        // `inserted_piece_indices` from already stored piece indices, leaving just those that are
510        // still missing in cache
511        caches.free_unneeded_stored_pieces(&mut piece_indices_to_store);
512
513        if let Some(metrics) = &self.metrics {
514            for offset in caches.stored_pieces_offsets() {
515                piece_caches_capacity_used[usize::from(offset.cache_index)] += 1;
516            }
517
518            for cache_used in piece_caches_capacity_used {
519                metrics
520                    .piece_cache_capacity_used
521                    .inc_by(i64::from(cache_used));
522            }
523        }
524
525        let stored_count = caches.stored_pieces_offsets().len();
526        // Store whatever correct pieces are immediately available after restart, the rest is
527        // updated in place as pieces are downloaded, so readers always see the latest state
528        *self.piece_caches.write().await = caches;
529
530        debug!(
531            %stored_count,
532            count = %piece_indices_to_store.len(),
533            "Identified piece indices that should be cached",
534        );
535
536        let pieces_to_download_total = piece_indices_to_store.len() + stored_count;
537        let piece_indices_to_store = piece_indices_to_store
538            .into_values()
539            .collect::<Vec<_>>()
540            // TODO: Allocating chunks here shouldn't be necessary, but otherwise it fails with
541            //  confusing error described in https://github.com/rust-lang/rust/issues/64552 and
542            //  similar upstream issues
543            .chunks(SYNC_BATCH_SIZE)
544            .map(<[PieceIndex]>::to_vec)
545            .collect::<Vec<_>>();
546
547        let downloaded_pieces_count = AtomicUsize::new(stored_count);
548        self.handlers.progress.call_simple(&0.0);
549        let batch_count = piece_indices_to_store.len();
550        let piece_indices_to_store = piece_indices_to_store.into_iter().enumerate();
551
552        let downloading_semaphore = &Semaphore::new(SYNC_BATCH_SIZE * SYNC_CONCURRENT_BATCHES);
553        let ignored_cache_indices = &RwLock::new(HashSet::new());
554
555        let downloading_pieces_stream =
556            stream::iter(piece_indices_to_store.map(|(batch, piece_indices)| {
557                let downloaded_pieces_count = &downloaded_pieces_count;
558                let num_pieces = piece_indices.len();
559
560                trace!(
561                    %num_pieces,
562                    %batch,
563                    %batch_count,
564                    first_piece_index = ?piece_indices.first().expect("chunks are never empty"),
565                    last_piece_index = ?piece_indices.last().expect("chunks are never empty"),
566                    downloaded_pieces_count = %downloaded_pieces_count.load(Ordering::Relaxed),
567                    %pieces_to_download_total,
568                    available_permits = %downloading_semaphore.available_permits(),
569                    "Started piece cache sync batch",
570                );
571
572                async move {
573                    let mut permit = downloading_semaphore
574                        .acquire_many(SYNC_BATCH_SIZE as u32)
575                        .await
576                        .expect("Semaphore is never closed; qed");
577                    debug!(%batch, %num_pieces, "Downloading pieces");
578
579                    let pieces_stream = match piece_getter.get_pieces(piece_indices).await {
580                        Ok(pieces_stream) => pieces_stream,
581                        Err(error) => {
582                            error!(
583                                %error,
584                                "Failed to get pieces from piece getter"
585                            );
586                            return;
587                        }
588                    };
589                    let mut pieces_stream = pieces_stream.enumerate();
590
591                    while let Some((index, (piece_index, result))) = pieces_stream.next().await {
592                        debug!(%batch, %index, %piece_index, "Downloaded piece");
593                        // Release slot for future batches, by dropping it along with the piece.
594                        let _permit = permit.split(1);
595
596                        let piece = match result {
597                            Ok(Some(piece)) => {
598                                trace!(%batch, %piece_index, "Downloaded piece successfully");
599                                piece
600                            }
601                            Ok(None) => {
602                                debug!(%batch, %piece_index, "Couldn't find piece");
603                                continue;
604                            }
605                            Err(error) => {
606                                debug!(
607                                    %batch,
608                                    %error,
609                                    %piece_index,
610                                    "Failed to get piece for piece cache"
611                                );
612                                continue;
613                            }
614                        };
615
616                        let (offset, maybe_backend) = {
617                            let mut caches = self.piece_caches.write().await;
618
619                            // Find plot in which there is a place for new piece to be stored
620                            let Some(offset) = caches.pop_free_offset() else {
621                                error!(
622                                    %batch,
623                                    %piece_index,
624                                    "Failed to store piece in cache, there was no space"
625                                );
626                                break;
627                            };
628
629                            (offset, caches.get_backend(offset.cache_index).cloned())
630                        };
631
632                        let cache_index = offset.cache_index;
633                        let piece_offset = offset.piece_offset;
634
635                        let skip_write = ignored_cache_indices.read().contains(&cache_index);
636                        if skip_write {
637                            trace!(
638                                %batch,
639                                %cache_index,
640                                %piece_index,
641                                %piece_offset,
642                                "Skipping known problematic cache index"
643                            );
644                        } else {
645                            if let Some(backend) = maybe_backend
646                                && let Err(error) =
647                                    backend.write_piece(piece_offset, piece_index, &piece).await
648                            {
649                                error!(
650                                    %error,
651                                    %batch,
652                                    %cache_index,
653                                    %piece_index,
654                                    %piece_offset,
655                                    "Failed to write piece into cache, ignoring this cache going \
656                                    forward"
657                                );
658                                ignored_cache_indices.write().insert(cache_index);
659                                continue;
660                            }
661
662                            let key =
663                                KeyWithDistance::new(self.peer_id, piece_index.to_multihash());
664                            self.piece_caches
665                                .write()
666                                .await
667                                .push_stored_piece(key, offset);
668                        }
669
670                        let prev_downloaded_pieces_count =
671                            downloaded_pieces_count.fetch_add(1, Ordering::Relaxed);
672                        // Do not print anything or send progress notification after last piece
673                        // until piece cache is written fully below
674                        if prev_downloaded_pieces_count != pieces_to_download_total {
675                            let progress = prev_downloaded_pieces_count as f32
676                                / pieces_to_download_total as f32
677                                * 100.0;
678                            if prev_downloaded_pieces_count
679                                .is_multiple_of(SYNC_PROGRESS_LOG_INTERVAL)
680                            {
681                                info!(
682                                    "Piece cache sync {progress:.2}% complete ({} / {})",
683                                    ByteSize::b(
684                                        (prev_downloaded_pieces_count * Piece::SIZE) as u64,
685                                    )
686                                    .display()
687                                    .iec(),
688                                    ByteSize::b((pieces_to_download_total * Piece::SIZE) as u64,)
689                                        .display()
690                                        .iec(),
691                                );
692                            }
693
694                            self.handlers.progress.call_simple(&progress);
695                        }
696                    }
697
698                    trace!(
699                        %num_pieces,
700                        %batch,
701                        %batch_count,
702                        downloaded_pieces_count = %downloaded_pieces_count.load(Ordering::Relaxed),
703                        %pieces_to_download_total,
704                        available_permits = %downloading_semaphore.available_permits(),
705                        "Finished piece cache sync batch",
706                    );
707                }
708            }));
709
710        // Download several batches concurrently to make sure slow tail of one is compensated by
711        // another
712        downloading_pieces_stream
713            // This allows to schedule new batch while previous batches partially completed, but
714            // avoids excessive memory usage like when all futures are created upfront
715            .buffer_unordered(SYNC_CONCURRENT_BATCHES * 10)
716            // Simply drain everything
717            .for_each(|()| async {})
718            .await;
719
720        self.handlers.progress.call_simple(&100.0);
721        *last_segment_index_internal = last_segment_index;
722
723        info!("Finished piece cache synchronization");
724    }
725
726    async fn process_new_segment_index<PG>(
727        &self,
728        piece_getter: &PG,
729        segment_index: SegmentIndex,
730        last_segment_index_internal: &mut SegmentIndex,
731    ) where
732        PG: PieceGetter,
733    {
734        debug!(
735            %last_segment_index_internal,
736            %segment_index,
737            "Starting to process new segment index"
738        );
739
740        let segment_index_range =
741            (*last_segment_index_internal + SegmentIndex::ONE)..=segment_index;
742
743        if segment_index_range.is_empty() {
744            return;
745        }
746
747        debug!(
748            ?segment_index_range,
749            "Downloading potentially useful pieces"
750        );
751
752        // TODO: Download local pieces first, then download other pieces only after some delay, so
753        //  other farmers on the network have enough time to store their pieces too
754
755        // We do not insert pieces into cache/heap yet, so we don't know if all of these pieces will
756        // be included, but there is a good chance they will be
757        let pieces_to_maybe_include = segment_index_range
758            .flat_map(|segment_index| segment_index.segment_piece_indexes())
759            .map(|piece_index| async move {
760                let should_store_in_piece_cache = self
761                    .piece_caches
762                    .read()
763                    .await
764                    .should_include_key(self.peer_id, piece_index);
765
766                let key = RecordKey::from(piece_index.to_multihash());
767                let should_store_in_plot_cache =
768                    self.plot_caches.should_store(piece_index, &key).await;
769
770                if !(should_store_in_piece_cache || should_store_in_plot_cache) {
771                    trace!(%piece_index, "Piece doesn't need to be cached #1");
772
773                    return None;
774                }
775
776                let maybe_piece_result =
777                    self.node_client
778                        .piece(piece_index)
779                        .await
780                        .inspect_err(|error| {
781                            debug!(
782                                %error,
783                                %segment_index,
784                                %piece_index,
785                                "Failed to retrieve piece from node right after archiving"
786                            );
787                        });
788
789                if let Ok(Some(piece)) = maybe_piece_result {
790                    return Some((piece_index, piece));
791                }
792
793                match piece_getter.get_piece(piece_index).await {
794                    Ok(Some(piece)) => Some((piece_index, piece)),
795                    Ok(None) => {
796                        warn!(
797                            %segment_index,
798                            %piece_index,
799                            "Failed to retrieve piece right after archiving"
800                        );
801
802                        None
803                    }
804                    Err(error) => {
805                        warn!(
806                            %error,
807                            %segment_index,
808                            %piece_index,
809                            "Failed to retrieve piece right after archiving"
810                        );
811
812                        None
813                    }
814                }
815            })
816            .collect::<FuturesUnordered<_>>()
817            .filter_map(|maybe_piece| async move { maybe_piece })
818            .collect::<Vec<_>>()
819            .await;
820
821        debug!(%segment_index, "Downloaded potentially useful pieces");
822
823        // Go through potentially matching pieces again now that segment was acknowledged and try to
824        // persist them if necessary
825        for (piece_index, piece) in pieces_to_maybe_include {
826            if !self
827                .plot_caches
828                .store_additional_piece(piece_index, &piece)
829                .await
830            {
831                trace!(%piece_index, "Piece could not be cached in plot cache");
832            }
833
834            if !self
835                .piece_caches
836                .read()
837                .await
838                .should_include_key(self.peer_id, piece_index)
839            {
840                trace!(%piece_index, "Piece doesn't need to be cached #2");
841
842                continue;
843            }
844
845            trace!(%piece_index, "Piece needs to be cached #1");
846
847            self.persist_piece_in_cache(piece_index, piece).await;
848        }
849
850        *last_segment_index_internal = segment_index;
851
852        debug!(%segment_index, "Finished processing new segment indices");
853    }
854
855    async fn keep_up_after_initial_sync<PG>(
856        &self,
857        piece_getter: &PG,
858        last_segment_index_internal: &mut SegmentIndex,
859    ) where
860        PG: PieceGetter,
861    {
862        let last_segment_index = match self.node_client.farmer_app_info().await {
863            Ok(farmer_app_info) => farmer_app_info.protocol_info.history_size.segment_index(),
864            Err(error) => {
865                error!(
866                    %error,
867                    "Failed to get farmer app info from node, keeping old cache state without \
868                    updates"
869                );
870                return;
871            }
872        };
873
874        if last_segment_index <= *last_segment_index_internal {
875            return;
876        }
877
878        info!(
879            "Syncing piece cache to the latest history size, this may pause block production if \
880            takes too long"
881        );
882
883        // Keep up with segment indices that were potentially created since reinitialization
884        let piece_indices = (*last_segment_index_internal..=last_segment_index)
885            .flat_map(|segment_index| segment_index.segment_piece_indexes());
886
887        // TODO: Download pieces concurrently
888        for piece_index in piece_indices {
889            if !self
890                .piece_caches
891                .read()
892                .await
893                .should_include_key(self.peer_id, piece_index)
894            {
895                trace!(%piece_index, "Piece doesn't need to be cached #3");
896
897                continue;
898            }
899
900            trace!(%piece_index, "Piece needs to be cached #2");
901
902            let result = piece_getter.get_piece(piece_index).await;
903
904            let piece = match result {
905                Ok(Some(piece)) => piece,
906                Ok(None) => {
907                    debug!(%piece_index, "Couldn't find piece");
908                    continue;
909                }
910                Err(error) => {
911                    debug!(
912                        %error,
913                        %piece_index,
914                        "Failed to get piece for piece cache"
915                    );
916                    continue;
917                }
918            };
919
920            self.persist_piece_in_cache(piece_index, piece).await;
921        }
922
923        info!("Finished syncing piece cache to the latest history size");
924
925        *last_segment_index_internal = last_segment_index;
926    }
927
928    /// This assumes it was already checked that piece needs to be stored, no verification for this
929    /// is done internally and invariants will break if this assumption doesn't hold true
930    async fn persist_piece_in_cache(&self, piece_index: PieceIndex, piece: Piece) {
931        let key = KeyWithDistance::new(self.peer_id, piece_index.to_multihash());
932        let mut caches = self.piece_caches.write().await;
933        if let Some((old_key, offset)) = caches.should_replace(&key) {
934            // Entry is already occupied, we need to find and replace old piece with new one
935            let cache_index = offset.cache_index;
936            let piece_offset = offset.piece_offset;
937            let Some(backend) = caches.get_backend(cache_index) else {
938                // Cache backend not exist
939                warn!(
940                    %cache_index,
941                    %piece_index,
942                    "Should have a cached backend, but it didn't exist, this is an \
943                    implementation bug"
944                );
945                return;
946            };
947            if let Err(error) = backend.write_piece(piece_offset, piece_index, &piece).await {
948                error!(
949                    %error,
950                    %cache_index,
951                    %piece_index,
952                    %piece_offset,
953                    "Failed to write piece into cache"
954                );
955            } else {
956                let old_piece_index = decode_piece_index_from_record_key(old_key.record_key());
957                trace!(
958                    %cache_index,
959                    %old_piece_index,
960                    %piece_index,
961                    %piece_offset,
962                    "Successfully replaced old cached piece"
963                );
964                caches.push_stored_piece(key, offset);
965            }
966        } else {
967            // There is free space in cache, need to find a free spot and place piece there
968
969            let Some(offset) = caches.pop_free_offset() else {
970                warn!(
971                    %piece_index,
972                    "Should have inserted piece into cache, but it didn't happen, this is an \
973                    implementation bug"
974                );
975                return;
976            };
977            let cache_index = offset.cache_index;
978            let piece_offset = offset.piece_offset;
979            let Some(backend) = caches.get_backend(cache_index) else {
980                // Cache backend not exist
981                warn!(
982                    %cache_index,
983                    %piece_index,
984                    "Should have a cached backend, but it didn't exist, this is an \
985                    implementation bug"
986                );
987                return;
988            };
989
990            if let Err(error) = backend.write_piece(piece_offset, piece_index, &piece).await {
991                error!(
992                    %error,
993                    %cache_index,
994                    %piece_index,
995                    %piece_offset,
996                    "Failed to write piece into cache"
997                );
998            } else {
999                trace!(
1000                    %cache_index,
1001                    %piece_index,
1002                    %piece_offset,
1003                    "Successfully stored piece in cache"
1004                );
1005                if let Some(metrics) = &self.metrics {
1006                    metrics.piece_cache_capacity_used.inc();
1007                }
1008                caches.push_stored_piece(key, offset);
1009            }
1010        }
1011    }
1012}
1013
1014#[derive(Debug)]
1015struct PlotCaches {
1016    /// Additional piece caches
1017    caches: AsyncRwLock<Vec<StdArc<dyn PlotCache>>>,
1018    /// Next plot cache to use for storing pieces
1019    next_plot_cache: AtomicUsize,
1020}
1021
1022impl PlotCaches {
1023    /// Returns true if there might be space to add a piece to a cache.
1024    /// Returns false if it is already in a cache, or it can't be added to any of the caches.
1025    ///
1026    /// Available space can be overwritten by a sector at any time, so the piece write can still
1027    /// fail even if this returns `true`.
1028    async fn should_store(&self, piece_index: PieceIndex, key: &RecordKey) -> bool {
1029        for (cache_index, cache) in self.caches.read().await.iter().enumerate() {
1030            match cache.is_piece_maybe_stored(key).await {
1031                Ok(MaybePieceStoredResult::No) => {
1032                    // Isn't stored or can't be stored, try another cache if there is one
1033                }
1034                Ok(MaybePieceStoredResult::Vacant) => {
1035                    return true;
1036                }
1037                Ok(MaybePieceStoredResult::Yes) => {
1038                    // Already stored, nothing else left to do
1039                    return false;
1040                }
1041                Err(error) => {
1042                    warn!(
1043                        %cache_index,
1044                        %piece_index,
1045                        %error,
1046                        "Failed to check piece stored in cache"
1047                    );
1048                }
1049            }
1050        }
1051
1052        false
1053    }
1054
1055    /// Store a piece in additional downloaded pieces, if there is space for it.
1056    /// Returns `true` if the piece was added to a cache, and `false` if it couldn't be stored,
1057    /// typically because the cache is full.
1058    async fn store_additional_piece(&self, piece_index: PieceIndex, piece: &Piece) -> bool {
1059        let plot_caches = self.caches.read().await;
1060        let plot_caches_len = plot_caches.len();
1061
1062        // Store pieces in plots using round-robin distribution
1063        for _ in 0..plot_caches_len {
1064            let plot_cache_index =
1065                self.next_plot_cache.fetch_add(1, Ordering::Relaxed) % plot_caches_len;
1066
1067            match plot_caches[plot_cache_index]
1068                .try_store_piece(piece_index, piece)
1069                .await
1070            {
1071                Ok(true) => {
1072                    return true;
1073                }
1074                Ok(false) => {
1075                    // Nothing
1076                }
1077                Err(error) => {
1078                    error!(
1079                        %error,
1080                        %piece_index,
1081                        %plot_cache_index,
1082                        "Failed to store additional piece in cache"
1083                    );
1084                }
1085            }
1086        }
1087
1088        false
1089    }
1090}
1091
1092/// Farmer cache that aggregates different kinds of caches of multiple disks.
1093///
1094/// Pieces in [`PieceCache`] are stored based on capacity and proximity of piece index to farmer's
1095/// network identity. If capacity is not enough to store all pieces in cache then pieces that are
1096/// further from network identity will be evicted, this is helpful for quick retrieval of pieces
1097/// from DSN as well as plotting purposes.
1098///
1099/// [`PlotCache`] is used as a supplementary cache and is primarily helpful for smaller farmers
1100/// where piece cache is not enough to store all the pieces on the network, while there is a lot of
1101/// space in the plot that is not used by sectors yet and can be leverage as extra caching space.
1102#[derive(Debug, Clone)]
1103pub struct FarmerCache {
1104    peer_id: PeerId,
1105    /// Individual dedicated piece caches
1106    piece_caches: Arc<AsyncRwLock<PieceCachesState>>,
1107    /// Additional piece caches
1108    plot_caches: Arc<PlotCaches>,
1109    handlers: Arc<Handlers>,
1110    // We do not want to increase capacity unnecessarily on clone
1111    worker_sender: mpsc::Sender<WorkerCommand>,
1112    metrics: Option<Arc<FarmerCacheMetrics>>,
1113}
1114
1115impl FarmerCache {
1116    /// Create new piece cache instance and corresponding worker.
1117    ///
1118    /// NOTE: Returned future is async, but does blocking operations and should be running in
1119    /// dedicated thread.
1120    pub fn new<NC>(
1121        node_client: NC,
1122        peer_id: PeerId,
1123        registry: Option<&mut Registry>,
1124    ) -> (Self, FarmerCacheWorker<NC>)
1125    where
1126        NC: NodeClient,
1127    {
1128        let caches = Arc::default();
1129        let (worker_sender, worker_receiver) = mpsc::channel(WORKER_CHANNEL_CAPACITY);
1130        let handlers = Arc::new(Handlers::default());
1131
1132        let plot_caches = Arc::new(PlotCaches {
1133            caches: AsyncRwLock::default(),
1134            next_plot_cache: AtomicUsize::new(0),
1135        });
1136        let metrics = registry.map(|registry| Arc::new(FarmerCacheMetrics::new(registry)));
1137
1138        let instance = Self {
1139            peer_id,
1140            piece_caches: Arc::clone(&caches),
1141            plot_caches: Arc::clone(&plot_caches),
1142            handlers: Arc::clone(&handlers),
1143            worker_sender,
1144            metrics: metrics.clone(),
1145        };
1146        let worker = FarmerCacheWorker {
1147            peer_id,
1148            node_client,
1149            piece_caches: caches,
1150            plot_caches,
1151            handlers,
1152            worker_receiver: Some(worker_receiver),
1153            metrics,
1154        };
1155
1156        (instance, worker)
1157    }
1158
1159    /// Get piece from cache
1160    pub async fn get_piece<Key>(&self, key: Key) -> Option<Piece>
1161    where
1162        RecordKey: From<Key>,
1163    {
1164        let key = RecordKey::from(key);
1165        let maybe_piece_found = {
1166            let key = KeyWithDistance::new_with_record_key(self.peer_id, key.clone());
1167            let caches = self.piece_caches.read().await;
1168
1169            caches.get_stored_piece(&key).and_then(|offset| {
1170                let cache_index = offset.cache_index;
1171                let piece_offset = offset.piece_offset;
1172                Some((
1173                    piece_offset,
1174                    cache_index,
1175                    caches.get_backend(cache_index)?.clone(),
1176                ))
1177            })
1178        };
1179
1180        if let Some((piece_offset, cache_index, backend)) = maybe_piece_found {
1181            match backend.read_piece(piece_offset).await {
1182                Ok(maybe_piece) => {
1183                    return if let Some((_piece_index, piece)) = maybe_piece {
1184                        if let Some(metrics) = &self.metrics {
1185                            metrics.cache_get_hit.inc();
1186                        }
1187                        Some(piece)
1188                    } else {
1189                        error!(
1190                            %cache_index,
1191                            %piece_offset,
1192                            ?key,
1193                            "Piece was expected to be in cache, but wasn't found there"
1194                        );
1195                        if let Some(metrics) = &self.metrics {
1196                            metrics.cache_get_error.inc();
1197                        }
1198                        None
1199                    };
1200                }
1201                Err(error) => {
1202                    error!(
1203                        %error,
1204                        %cache_index,
1205                        ?key,
1206                        %piece_offset,
1207                        "Error while reading piece from cache"
1208                    );
1209
1210                    if let Err(error) = self
1211                        .worker_sender
1212                        .clone()
1213                        .send(WorkerCommand::ForgetKey { key })
1214                        .await
1215                    {
1216                        trace!(%error, "Failed to send ForgetKey command to worker");
1217                    }
1218
1219                    if let Some(metrics) = &self.metrics {
1220                        metrics.cache_get_error.inc();
1221                    }
1222                    return None;
1223                }
1224            }
1225        }
1226
1227        for cache in self.plot_caches.caches.read().await.iter() {
1228            if let Ok(Some(piece)) = cache.read_piece(&key).await {
1229                if let Some(metrics) = &self.metrics {
1230                    metrics.cache_get_hit.inc();
1231                }
1232                return Some(piece);
1233            }
1234        }
1235
1236        if let Some(metrics) = &self.metrics {
1237            metrics.cache_get_miss.inc();
1238        }
1239        None
1240    }
1241
1242    /// Get pieces from cache.
1243    ///
1244    /// Number of elements in returned stream is the same as number of unique `piece_indices`.
1245    pub async fn get_pieces<'a, PieceIndices>(
1246        &'a self,
1247        piece_indices: PieceIndices,
1248    ) -> impl Stream<Item = (PieceIndex, Option<Piece>)> + Send + Unpin + 'a
1249    where
1250        PieceIndices: IntoIterator<Item = PieceIndex, IntoIter: Send + 'a> + Send + 'a,
1251    {
1252        let mut pieces_to_get_from_plot_cache = Vec::new();
1253
1254        let pieces_to_read_from_piece_cache = {
1255            let caches = self.piece_caches.read().await;
1256            // Pieces to read from piece cache grouped by backend for efficiency reasons
1257            let mut pieces_to_read_from_piece_cache =
1258                HashMap::<CacheIndex, (CacheBackend, HashMap<_, _>)>::new();
1259
1260            for piece_index in piece_indices {
1261                let key = RecordKey::from(piece_index.to_multihash());
1262
1263                let Some(offset) = caches.get_stored_piece(&KeyWithDistance::new_with_record_key(
1264                    self.peer_id,
1265                    key.clone(),
1266                )) else {
1267                    pieces_to_get_from_plot_cache.push((piece_index, key));
1268                    continue;
1269                };
1270
1271                let cache_index = offset.cache_index;
1272                let piece_offset = offset.piece_offset;
1273
1274                match pieces_to_read_from_piece_cache.entry(cache_index) {
1275                    Entry::Occupied(mut entry) => {
1276                        let (_backend, pieces) = entry.get_mut();
1277                        pieces.insert(piece_offset, (piece_index, key));
1278                    }
1279                    Entry::Vacant(entry) => {
1280                        let backend = if let Some(backend) = caches.get_backend(cache_index) {
1281                            backend.clone()
1282                        } else {
1283                            pieces_to_get_from_plot_cache.push((piece_index, key));
1284                            continue;
1285                        };
1286                        entry
1287                            .insert((backend, HashMap::from([(piece_offset, (piece_index, key))])));
1288                    }
1289                }
1290            }
1291
1292            pieces_to_read_from_piece_cache
1293        };
1294
1295        let (tx, mut rx) = mpsc::unbounded();
1296
1297        let fut = async move {
1298            let tx = &tx;
1299
1300            let mut reading_from_piece_cache = pieces_to_read_from_piece_cache
1301                .into_iter()
1302                .map(|(cache_index, (backend, mut pieces_to_get))| async move {
1303                    let mut pieces_stream = match backend
1304                        .read_pieces(Box::new(
1305                            pieces_to_get
1306                                .keys()
1307                                .copied()
1308                                .collect::<Vec<_>>()
1309                                .into_iter(),
1310                        ))
1311                        .await
1312                    {
1313                        Ok(pieces_stream) => pieces_stream,
1314                        Err(error) => {
1315                            error!(
1316                                %error,
1317                                %cache_index,
1318                                "Error while reading pieces from cache"
1319                            );
1320
1321                            if let Some(metrics) = &self.metrics {
1322                                metrics.cache_get_error.inc_by(pieces_to_get.len() as u64);
1323                            }
1324                            for (piece_index, _key) in pieces_to_get.into_values() {
1325                                tx.unbounded_send((piece_index, None)).expect(
1326                                    "This future isn't polled after receiver is dropped; qed",
1327                                );
1328                            }
1329                            return;
1330                        }
1331                    };
1332
1333                    while let Some(maybe_piece) = pieces_stream.next().await {
1334                        let result = match maybe_piece {
1335                            Ok((piece_offset, Some((piece_index, piece)))) => {
1336                                pieces_to_get.remove(&piece_offset);
1337
1338                                if let Some(metrics) = &self.metrics {
1339                                    metrics.cache_get_hit.inc();
1340                                }
1341                                (piece_index, Some(piece))
1342                            }
1343                            Ok((piece_offset, None)) => {
1344                                let Some((piece_index, key)) = pieces_to_get.remove(&piece_offset)
1345                                else {
1346                                    debug!(
1347                                        %cache_index,
1348                                        %piece_offset,
1349                                        "Received piece offset that was not expected"
1350                                    );
1351                                    continue;
1352                                };
1353
1354                                error!(
1355                                    %cache_index,
1356                                    %piece_index,
1357                                    %piece_offset,
1358                                    ?key,
1359                                    "Piece was expected to be in cache, but wasn't found there"
1360                                );
1361                                if let Some(metrics) = &self.metrics {
1362                                    metrics.cache_get_error.inc();
1363                                }
1364                                (piece_index, None)
1365                            }
1366                            Err(error) => {
1367                                error!(
1368                                    %error,
1369                                    %cache_index,
1370                                    "Error while reading piece from cache"
1371                                );
1372
1373                                if let Some(metrics) = &self.metrics {
1374                                    metrics.cache_get_error.inc();
1375                                }
1376                                continue;
1377                            }
1378                        };
1379
1380                        tx.unbounded_send(result)
1381                            .expect("This future isn't polled after receiver is dropped; qed");
1382                    }
1383
1384                    if pieces_to_get.is_empty() {
1385                        return;
1386                    }
1387
1388                    if let Some(metrics) = &self.metrics {
1389                        metrics.cache_get_error.inc_by(pieces_to_get.len() as u64);
1390                    }
1391                    for (piece_offset, (piece_index, key)) in pieces_to_get {
1392                        error!(
1393                            %cache_index,
1394                            %piece_index,
1395                            %piece_offset,
1396                            ?key,
1397                            "Piece cache didn't return an entry for offset"
1398                        );
1399
1400                        // Uphold invariant of the method that some result should be returned
1401                        // for every unique piece index
1402                        tx.unbounded_send((piece_index, None))
1403                            .expect("This future isn't polled after receiver is dropped; qed");
1404                    }
1405                })
1406                .collect::<FuturesUnordered<_>>();
1407            // TODO: Can't use this due to https://github.com/rust-lang/rust/issues/64650
1408            // Simply drain everything
1409            // .for_each(|()| async {})
1410
1411            // TODO: Remove once https://github.com/rust-lang/rust/issues/64650 is resolved
1412            let reading_from_piece_cache_fut = async move {
1413                while let Some(()) = reading_from_piece_cache.next().await {
1414                    // Simply drain everything
1415                }
1416            };
1417
1418            let reading_from_plot_cache_fut = async {
1419                if pieces_to_get_from_plot_cache.is_empty() {
1420                    return;
1421                }
1422
1423                for cache in self.plot_caches.caches.read().await.iter() {
1424                    // Iterating over offsets in reverse order to both traverse elements in async
1425                    // code and being able to efficiently remove entries without extra allocations
1426                    for offset in (0..pieces_to_get_from_plot_cache.len()).rev() {
1427                        let (piece_index, key) = &pieces_to_get_from_plot_cache[offset];
1428
1429                        if let Ok(Some(piece)) = cache.read_piece(key).await {
1430                            if let Some(metrics) = &self.metrics {
1431                                metrics.cache_get_hit.inc();
1432                            }
1433                            tx.unbounded_send((*piece_index, Some(piece)))
1434                                .expect("This future isn't polled after receiver is dropped; qed");
1435
1436                            // Due to iteration in reverse order and swapping using elements at the
1437                            // end, this doesn't affect processing of the elements
1438                            pieces_to_get_from_plot_cache.swap_remove(offset);
1439                        }
1440                    }
1441
1442                    if pieces_to_get_from_plot_cache.is_empty() {
1443                        return;
1444                    }
1445                }
1446
1447                if let Some(metrics) = &self.metrics {
1448                    metrics
1449                        .cache_get_miss
1450                        .inc_by(pieces_to_get_from_plot_cache.len() as u64);
1451                }
1452
1453                for (piece_index, _key) in pieces_to_get_from_plot_cache {
1454                    tx.unbounded_send((piece_index, None))
1455                        .expect("This future isn't polled after receiver is dropped; qed");
1456                }
1457            };
1458
1459            join!(reading_from_piece_cache_fut, reading_from_plot_cache_fut).await
1460        };
1461        let mut fut = Box::pin(fut.fuse());
1462
1463        // Drive above future and stream back any pieces that were downloaded so far
1464        stream::poll_fn(move |cx| {
1465            if !fut.is_terminated() {
1466                // Result doesn't matter, we'll need to poll stream below anyway
1467                let _: Poll<_> = fut.poll_unpin(cx);
1468            }
1469
1470            if let Poll::Ready(maybe_result) = rx.poll_next_unpin(cx) {
1471                return Poll::Ready(maybe_result);
1472            }
1473
1474            // Exit will be done by the stream above
1475            Poll::Pending
1476        })
1477    }
1478
1479    /// Returns a filtered list of pieces that were found in farmer cache, order is not guaranteed
1480    pub async fn has_pieces(&self, mut piece_indices: Vec<PieceIndex>) -> Vec<PieceIndex> {
1481        let mut pieces_to_find = HashMap::<PieceIndex, RecordKey>::from_iter(
1482            piece_indices
1483                .iter()
1484                .map(|piece_index| (*piece_index, RecordKey::from(piece_index.to_multihash()))),
1485        );
1486
1487        // Quick check in piece caches
1488        {
1489            let piece_caches = self.piece_caches.read().await;
1490            pieces_to_find.retain(|_piece_index, key| {
1491                let distance_key = KeyWithDistance::new(self.peer_id, key.clone());
1492                !piece_caches.contains_stored_piece(&distance_key)
1493            });
1494        }
1495
1496        // Early exit if everything is cached
1497        if pieces_to_find.is_empty() {
1498            return piece_indices;
1499        }
1500
1501        // Check plot caches concurrently
1502        if let Some(plot_caches) = self.plot_caches.caches.try_read() {
1503            let plot_caches = &plot_caches;
1504            let not_found = pieces_to_find
1505                .into_iter()
1506                .map(|(piece_index, key)| async move {
1507                    let key = &key;
1508
1509                    let found = plot_caches
1510                        .iter()
1511                        .map(|plot_cache| async {
1512                            matches!(
1513                                plot_cache.is_piece_maybe_stored(key).await,
1514                                Ok(MaybePieceStoredResult::Yes)
1515                            )
1516                        })
1517                        .collect::<FuturesUnordered<_>>()
1518                        .any(|found| async move { found })
1519                        .await;
1520
1521                    if found { None } else { Some(piece_index) }
1522                })
1523                .collect::<FuturesUnordered<_>>()
1524                .filter_map(|maybe_piece_index| async move { maybe_piece_index })
1525                .collect::<HashSet<_>>()
1526                .await;
1527            piece_indices.retain(|piece_index| !not_found.contains(piece_index));
1528        }
1529        piece_indices
1530    }
1531
1532    /// Find piece in cache and return its retrieval details
1533    pub async fn find_piece(
1534        &self,
1535        piece_index: PieceIndex,
1536    ) -> Option<(PieceCacheId, PieceCacheOffset)> {
1537        let caches = self.piece_caches.read().await;
1538
1539        self.find_piece_internal(&caches, piece_index)
1540    }
1541
1542    /// Find pieces in cache and return their retrieval details
1543    pub async fn find_pieces<PieceIndices>(
1544        &self,
1545        piece_indices: PieceIndices,
1546    ) -> Vec<(PieceIndex, PieceCacheId, PieceCacheOffset)>
1547    where
1548        PieceIndices: IntoIterator<Item = PieceIndex>,
1549    {
1550        let caches = self.piece_caches.read().await;
1551
1552        piece_indices
1553            .into_iter()
1554            .filter_map(|piece_index| {
1555                self.find_piece_internal(&caches, piece_index)
1556                    .map(|(cache_id, piece_offset)| (piece_index, cache_id, piece_offset))
1557            })
1558            .collect()
1559    }
1560
1561    fn find_piece_internal(
1562        &self,
1563        caches: &PieceCachesState,
1564        piece_index: PieceIndex,
1565    ) -> Option<(PieceCacheId, PieceCacheOffset)> {
1566        let key = KeyWithDistance::new(self.peer_id, piece_index.to_multihash());
1567
1568        let Some(offset) = caches.get_stored_piece(&key) else {
1569            if let Some(metrics) = &self.metrics {
1570                metrics.cache_find_miss.inc();
1571            }
1572
1573            return None;
1574        };
1575        let piece_offset = offset.piece_offset;
1576
1577        if let Some(backend) = caches.get_backend(offset.cache_index) {
1578            if let Some(metrics) = &self.metrics {
1579                metrics.cache_find_hit.inc();
1580            }
1581            return Some((*backend.id(), piece_offset));
1582        }
1583
1584        if let Some(metrics) = &self.metrics {
1585            metrics.cache_find_miss.inc();
1586        }
1587        None
1588    }
1589
1590    /// Try to store a piece in additional downloaded pieces, if there is space for it.
1591    /// Returns `true` if the piece was added to this cache, and `false` if it was already stored,
1592    /// or there was no space.
1593    pub async fn maybe_store_additional_piece(
1594        &self,
1595        piece_index: PieceIndex,
1596        piece: &Piece,
1597    ) -> bool {
1598        let key = RecordKey::from(piece_index.to_multihash());
1599
1600        let should_store = self.plot_caches.should_store(piece_index, &key).await;
1601
1602        if !should_store {
1603            return false;
1604        }
1605
1606        self.plot_caches
1607            .store_additional_piece(piece_index, piece)
1608            .await
1609    }
1610
1611    /// Initialize replacement of backing caches
1612    pub async fn replace_backing_caches(
1613        &self,
1614        new_piece_caches: Vec<StdArc<dyn PieceCache>>,
1615        new_plot_caches: Vec<StdArc<dyn PlotCache>>,
1616    ) {
1617        if let Err(error) = self
1618            .worker_sender
1619            .clone()
1620            .send(WorkerCommand::ReplaceBackingCaches { new_piece_caches })
1621            .await
1622        {
1623            warn!(%error, "Failed to replace backing caches, worker exited");
1624        }
1625
1626        *self.plot_caches.caches.write().await = new_plot_caches;
1627    }
1628
1629    /// Subscribe to cache sync notifications
1630    pub fn on_sync_progress(&self, callback: HandlerFn<f32>) -> HandlerId {
1631        self.handlers.progress.add(callback)
1632    }
1633}
1634
1635/// Collection of [`FarmerCache`] instances for load balancing
1636#[derive(Debug, Clone)]
1637pub struct FarmerCaches {
1638    caches: Arc<[FarmerCache]>,
1639}
1640
1641impl From<Arc<[FarmerCache]>> for FarmerCaches {
1642    fn from(caches: Arc<[FarmerCache]>) -> Self {
1643        Self { caches }
1644    }
1645}
1646
1647impl From<FarmerCache> for FarmerCaches {
1648    fn from(cache: FarmerCache) -> Self {
1649        Self {
1650            caches: Arc::from_iter([cache]),
1651        }
1652    }
1653}
1654
1655impl FarmerCaches {
1656    /// Get piece from cache
1657    pub async fn get_piece<Key>(&self, key: Key) -> Option<Piece>
1658    where
1659        RecordKey: From<Key>,
1660    {
1661        let farmer_cache = self.caches.choose(&mut rand::rng())?;
1662        farmer_cache.get_piece(key).await
1663    }
1664
1665    /// Get pieces from cache.
1666    ///
1667    /// Number of elements in returned stream is the same as number of unique `piece_indices`.
1668    pub async fn get_pieces<'a, PieceIndices>(
1669        &'a self,
1670        piece_indices: PieceIndices,
1671    ) -> impl Stream<Item = (PieceIndex, Option<Piece>)> + Send + Unpin + 'a
1672    where
1673        PieceIndices: IntoIterator<Item = PieceIndex, IntoIter: Send + 'a> + Send + 'a,
1674    {
1675        let Some(farmer_cache) = self.caches.choose(&mut rand::rng()) else {
1676            return Either::Left(stream::iter(
1677                piece_indices
1678                    .into_iter()
1679                    .map(|piece_index| (piece_index, None)),
1680            ));
1681        };
1682
1683        Either::Right(farmer_cache.get_pieces(piece_indices).await)
1684    }
1685
1686    /// Returns a filtered list of pieces that were found in farmer cache, order is not guaranteed
1687    pub async fn has_pieces(&self, piece_indices: Vec<PieceIndex>) -> Vec<PieceIndex> {
1688        let Some(farmer_cache) = self.caches.choose(&mut rand::rng()) else {
1689            return Vec::new();
1690        };
1691
1692        farmer_cache.has_pieces(piece_indices).await
1693    }
1694
1695    /// Try to store a piece in additional downloaded pieces, if there is space for it.
1696    /// Returns `true` if the piece was added to one or more caches, and `false` if it was already
1697    /// stored, or there was no space.
1698    pub async fn maybe_store_additional_piece(
1699        &self,
1700        piece_index: PieceIndex,
1701        piece: &Piece,
1702    ) -> bool {
1703        // Run all the futures to completion, and take a non-short-circuiting any() on the results.
1704        self.caches
1705            .iter()
1706            .map(|farmer_cache| farmer_cache.maybe_store_additional_piece(piece_index, piece))
1707            .collect::<FuturesUnordered<_>>()
1708            .fold::<bool, _, _>(false, |acc, stored| async move { acc || stored })
1709            .await
1710    }
1711}
1712
1713/// Extracts the `PieceIndex` from a `RecordKey`.
1714fn decode_piece_index_from_record_key(key: &RecordKey) -> PieceIndex {
1715    let len = key.as_ref().len();
1716    let s = len - PieceIndex::SIZE;
1717
1718    let mut piece_index_bytes = [0u8; PieceIndex::SIZE];
1719    piece_index_bytes.copy_from_slice(&key.as_ref()[s..]);
1720
1721    PieceIndex::from_bytes(piece_index_bytes)
1722}