Merge pull request #2550 from subspace/farmer-plot-cache

Farmer plot cache

Author: nazar-pc

crates/subspace-farmer/src/bin/subspace-farmer/commands/farm.rs

@@ -631,6 +631,10 @@ where
                 .iter()
                 .map(|single_disk_farm| single_disk_farm.piece_cache())
                 .collect(),
+            single_disk_farms
+                .iter()
+                .map(|single_disk_farm| single_disk_farm.plot_cache())
+                .collect(),
         )
         .await;
     drop(farmer_cache);

crates/subspace-farmer/src/farmer_cache.rs

@@ -3,6 +3,7 @@ mod tests;
 
 use crate::node_client::NodeClient;
 use crate::single_disk_farm::piece_cache::{DiskPieceCache, Offset};
+use crate::single_disk_farm::plot_cache::{DiskPlotCache, MaybePieceStoredResult};
 use crate::utils::{run_future_in_dedicated_thread, AsyncJoinOnDrop};
 use event_listener_primitives::{Bag, HandlerId};
 use futures::channel::oneshot;
@@ -11,6 +12,7 @@ use futures::{select, FutureExt, StreamExt};
 use parking_lot::RwLock;
 use std::collections::{HashMap, VecDeque};
 use std::num::NonZeroU16;
+use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
 use std::{fmt, mem};
@@ -51,7 +53,7 @@ struct DiskPieceCacheState {
 #[derive(Debug)]
 enum WorkerCommand {
     ReplaceBackingCaches {
-        new_caches: Vec<DiskPieceCache>,
+        new_piece_caches: Vec<DiskPieceCache>,
         acknowledgement: oneshot::Sender<()>,
     },
     ForgetKey {
@@ -102,11 +104,11 @@ where
             .expect("Always set during worker instantiation");
 
         if let Some(WorkerCommand::ReplaceBackingCaches {
-            new_caches,
+            new_piece_caches,
             acknowledgement,
         }) = worker_receiver.recv().await
         {
-            self.initialize(&piece_getter, &mut worker_state, new_caches)
+            self.initialize(&piece_getter, &mut worker_state, new_piece_caches)
                 .await;
             // Doesn't matter if receiver is still waiting for acknowledgement
             let _ = acknowledgement.send(());
@@ -163,10 +165,10 @@ where
     {
         match command {
             WorkerCommand::ReplaceBackingCaches {
-                new_caches,
+                new_piece_caches,
                 acknowledgement,
             } => {
-                self.initialize(piece_getter, worker_state, new_caches)
+                self.initialize(piece_getter, worker_state, new_piece_caches)
                     .await;
                 // Doesn't matter if receiver is still waiting for acknowledgement
                 let _ = acknowledgement.send(());
@@ -215,31 +217,31 @@ where
         &self,
         piece_getter: &PG,
         worker_state: &mut CacheWorkerState,
-        new_caches: Vec<DiskPieceCache>,
+        new_piece_caches: Vec<DiskPieceCache>,
     ) where
         PG: PieceGetter,
     {
         info!("Initializing piece cache");
         // Pull old cache state since it will be replaced with a new one and reuse its allocations
         let cache_state = mem::take(&mut *self.caches.write());
-        let mut stored_pieces = Vec::with_capacity(new_caches.len());
-        let mut free_offsets = Vec::with_capacity(new_caches.len());
+        let mut stored_pieces = Vec::with_capacity(new_piece_caches.len());
+        let mut free_offsets = Vec::with_capacity(new_piece_caches.len());
         for mut state in cache_state {
             state.stored_pieces.clear();
             stored_pieces.push(state.stored_pieces);
             state.free_offsets.clear();
             free_offsets.push(state.free_offsets);
         }
-        stored_pieces.resize(new_caches.len(), HashMap::default());
-        free_offsets.resize(new_caches.len(), VecDeque::default());
+        stored_pieces.resize(new_piece_caches.len(), HashMap::default());
+        free_offsets.resize(new_piece_caches.len(), VecDeque::default());
 
         debug!("Collecting pieces that were in the cache before");
 
         // Build cache state of all backends
         let maybe_caches_futures = stored_pieces
             .into_iter()
             .zip(free_offsets)
-            .zip(new_caches)
+            .zip(new_piece_caches)
             .enumerate()
             .map(
                 |(index, ((mut stored_pieces, mut free_offsets), new_cache))| {
@@ -760,8 +762,12 @@ where
 #[derive(Debug, Clone)]
 pub struct FarmerCache {
     peer_id: PeerId,
-    /// Individual disk caches where pieces are stored
-    caches: Arc<RwLock<Vec<DiskPieceCacheState>>>,
+    /// Individual dedicated piece caches
+    piece_caches: Arc<RwLock<Vec<DiskPieceCacheState>>>,
+    /// Additional piece caches
+    plot_caches: Arc<RwLock<Vec<DiskPlotCache>>>,
+    /// Next plot cache to use for storing pieces
+    next_plot_cache: Arc<AtomicUsize>,
     handlers: Arc<Handlers>,
     // We do not want to increase capacity unnecessarily on clone
     worker_sender: Arc<mpsc::Sender<WorkerCommand>>,
@@ -782,7 +788,9 @@ impl FarmerCache {
 
         let instance = Self {
             peer_id,
-            caches: Arc::clone(&caches),
+            piece_caches: Arc::clone(&caches),
+            plot_caches: Arc::default(),
+            next_plot_cache: Arc::new(AtomicUsize::new(0)),
             handlers: Arc::clone(&handlers),
             worker_sender: Arc::new(worker_sender),
         };
@@ -801,34 +809,47 @@ impl FarmerCache {
     pub async fn get_piece(&self, key: RecordKey) -> Option<Piece> {
         let maybe_piece_fut = tokio::task::spawn_blocking({
             let key = key.clone();
-            let caches = Arc::clone(&self.caches);
+            let piece_caches = Arc::clone(&self.piece_caches);
+            let plot_caches = Arc::clone(&self.plot_caches);
             let worker_sender = Arc::clone(&self.worker_sender);
 
             move || {
-                for (disk_farm_index, cache) in caches.read().iter().enumerate() {
-                    let Some(&offset) = cache.stored_pieces.get(&key) else {
-                        continue;
-                    };
-                    match cache.backend.read_piece(offset) {
-                        Ok(maybe_piece) => {
-                            return maybe_piece;
-                        }
-                        Err(error) => {
-                            error!(
-                                %error,
-                                %disk_farm_index,
-                                ?key,
-                                %offset,
-                                "Error while reading piece from cache, might be a disk corruption"
-                            );
+                {
+                    let piece_caches = piece_caches.read();
+                    for (disk_farm_index, cache) in piece_caches.iter().enumerate() {
+                        let Some(&offset) = cache.stored_pieces.get(&key) else {
+                            continue;
+                        };
+                        match cache.backend.read_piece(offset) {
+                            Ok(maybe_piece) => {
+                                return maybe_piece;
+                            }
+                            Err(error) => {
+                                error!(
+                                    %error,
+                                    %disk_farm_index,
+                                    ?key,
+                                    %offset,
+                                    "Error while reading piece from cache, might be a disk corruption"
+                                );
+
+                                if let Err(error) =
+                                    worker_sender.blocking_send(WorkerCommand::ForgetKey { key })
+                                {
+                                    trace!(%error, "Failed to send ForgetKey command to worker");
+                                }
 
-                            if let Err(error) =
-                                worker_sender.blocking_send(WorkerCommand::ForgetKey { key })
-                            {
-                                trace!(%error, "Failed to send ForgetKey command to worker");
+                                return None;
                             }
+                        }
+                    }
+                }
 
-                            return None;
+                {
+                    let plot_caches = plot_caches.read();
+                    for cache in plot_caches.iter() {
+                        if let Some(piece) = cache.read_piece(&key) {
+                            return Some(piece);
                         }
                     }
                 }
@@ -846,24 +867,92 @@ impl FarmerCache {
         }
     }
 
+    /// Try to store a piece in additional downloaded pieces, if there is space for them
+    pub async fn maybe_store_additional_piece(&self, piece_index: PieceIndex, piece: &Piece) {
+        let key = RecordKey::from(piece_index.to_multihash());
+
+        let mut should_store = false;
+        for cache in self.plot_caches.read().iter() {
+            match cache.is_piece_maybe_stored(&key) {
+                MaybePieceStoredResult::No => {
+                    // Try another one if there is any
+                }
+                MaybePieceStoredResult::Vacant => {
+                    should_store = true;
+                    break;
+                }
+                MaybePieceStoredResult::Yes => {
+                    // Already stored, nothing else left to do
+                    return;
+                }
+            }
+        }
+
+        if !should_store {
+            return;
+        }
+
+        let should_store_fut = tokio::task::spawn_blocking({
+            let plot_caches = Arc::clone(&self.plot_caches);
+            let next_plot_cache = Arc::clone(&self.next_plot_cache);
+            let piece = piece.clone();
+
+            move || {
+                let plot_caches = plot_caches.read();
+                let plot_caches_len = plot_caches.len();
+
+                // Store pieces in plots using round-robin distribution
+                for _ in 0..plot_caches_len {
+                    let plot_cache_index =
+                        next_plot_cache.fetch_add(1, Ordering::Relaxed) % plot_caches_len;
+
+                    match plot_caches[plot_cache_index].try_store_piece(piece_index, &piece) {
+                        Ok(true) => {
+                            return;
+                        }
+                        Ok(false) => {
+                            continue;
+                        }
+                        Err(error) => {
+                            error!(
+                                %error,
+                                %piece_index,
+                                %plot_cache_index,
+                                "Failed to store additional piece in cache"
+                            );
+                            continue;
+                        }
+                    }
+                }
+            }
+        });
+
+        if let Err(error) = AsyncJoinOnDrop::new(should_store_fut, true).await {
+            error!(%error, %piece_index, "Failed to store additional piece in cache");
+        }
+    }
+
     /// Initialize replacement of backing caches, returns acknowledgement receiver that can be used
     /// to identify when cache initialization has finished
     pub async fn replace_backing_caches(
         &self,
-        new_caches: Vec<DiskPieceCache>,
+        new_piece_caches: Vec<DiskPieceCache>,
+        new_plot_caches: Vec<DiskPlotCache>,
     ) -> oneshot::Receiver<()> {
         let (sender, receiver) = oneshot::channel();
         if let Err(error) = self
             .worker_sender
             .send(WorkerCommand::ReplaceBackingCaches {
-                new_caches,
+                new_piece_caches,
                 acknowledgement: sender,
             })
             .await
         {
             warn!(%error, "Failed to replace backing caches, worker exited");
         }
 
+        *self.plot_caches.write() = new_plot_caches;
+
         receiver
     }
 
@@ -876,10 +965,27 @@ impl FarmerCache {
 impl LocalRecordProvider for FarmerCache {
     fn record(&self, key: &RecordKey) -> Option<ProviderRecord> {
         // It is okay to take read lock here, writes locks are very infrequent and very short
-        for cache in self.caches.read().iter() {
-            if cache.stored_pieces.contains_key(key) {
+        for piece_cache in self.piece_caches.read().iter() {
+            if piece_cache.stored_pieces.contains_key(key) {
+                // Note: We store our own provider records locally without local addresses
+                // to avoid redundant storage and outdated addresses. Instead, these are
+                // acquired on demand when returning a `ProviderRecord` for the local node.
+                return Some(ProviderRecord {
+                    key: key.clone(),
+                    provider: self.peer_id,
+                    expires: None,
+                    addresses: Vec::new(),
+                });
+            };
+        }
+        // It is okay to take read lock here, writes locks almost never happen
+        for plot_cache in self.plot_caches.read().iter() {
+            if matches!(
+                plot_cache.is_piece_maybe_stored(key),
+                MaybePieceStoredResult::Yes
+            ) {
                 // Note: We store our own provider records locally without local addresses
-                // to avoid redundant storage and outdated addresses. Instead these are
+                // to avoid redundant storage and outdated addresses. Instead, these are
                 // acquired on demand when returning a `ProviderRecord` for the local node.
                 return Some(ProviderRecord {
                     key: key.clone(),

crates/subspace-farmer/src/farmer_cache/tests.rs

@@ -192,10 +192,13 @@ async fn basic() {
             tokio::spawn(farmer_cache_worker.run(piece_getter.clone()));
 
         let initialized_fut = farmer_cache
-            .replace_backing_caches(vec![
-                DiskPieceCache::open(path1.as_ref(), 1).unwrap(),
-                DiskPieceCache::open(path2.as_ref(), 1).unwrap(),
-            ])
+            .replace_backing_caches(
+                vec![
+                    DiskPieceCache::open(path1.as_ref(), 1).unwrap(),
+                    DiskPieceCache::open(path2.as_ref(), 1).unwrap(),
+                ],
+                vec![],
+            )
             .await;
 
         // Wait for piece cache to be initialized
@@ -375,10 +378,13 @@ async fn basic() {
 
         // Reopen with the same backing caches
         let initialized_fut = farmer_cache
-            .replace_backing_caches(vec![
-                DiskPieceCache::open(path1.as_ref(), 1).unwrap(),
-                DiskPieceCache::open(path2.as_ref(), 1).unwrap(),
-            ])
+            .replace_backing_caches(
+                vec![
+                    DiskPieceCache::open(path1.as_ref(), 1).unwrap(),
+                    DiskPieceCache::open(path2.as_ref(), 1).unwrap(),
+                ],
+                vec![],
+            )
             .await;
         drop(farmer_cache);