Merge pull request #992 from subspace/dsn-node-publish-piece-batcher

Node. Add piece publishing batcher.

Author: nazar-pc

Cargo.lock

@@ -35,6 +35,7 @@ futures = "0.3.25"
 rand = { version = "0.8.5", features = ["min_const_gen"] }
 subspace-archiving = { version = "0.1.0", path = "../subspace-archiving"}
 subspace-farmer-components = { version = "0.1.0", path = "../subspace-farmer-components" }
+tokio = { version = "1.21.2", features = ["sync"] }
 
 [features]
 default = ["std"]

crates/sp-lightclient/Cargo.toml

@@ -107,6 +107,7 @@ impl Farmer {
         };
         let sector_codec = SectorCodec::new(PLOT_SECTOR_SIZE as usize).unwrap();
         let piece_receiver_batch_size = 20usize;
+
         block_on(plot_sector(
             &public_key,
             sector_index,
@@ -117,7 +118,7 @@ impl Farmer {
             &sector_codec,
             Cursor::new(sector.as_mut_slice()),
             Cursor::new(sector_metadata.as_mut_slice()),
-            piece_receiver_batch_size,
+            &tokio::sync::Semaphore::new(piece_receiver_batch_size),
         ))
         .unwrap();
 

crates/sp-lightclient/src/tests.rs

@@ -19,6 +19,7 @@ use subspace_farmer_components::farming::audit_sector;
 use subspace_farmer_components::file_ext::FileExt;
 use subspace_farmer_components::plotting::plot_sector;
 use subspace_farmer_components::FarmerProtocolInfo;
+use tokio::sync::Semaphore;
 use utils::BenchPieceReceiver;
 
 mod utils;
@@ -64,6 +65,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     let global_challenge = Blake2b256Hash::default();
     let solution_range = SolutionRange::MAX;
     let piece_receiver_batch_size = 20usize;
+    let piece_receiver_semaphore = Semaphore::new(piece_receiver_batch_size);
 
     let plotted_sector = {
         let mut plotted_sector = vec![0u8; PLOT_SECTOR_SIZE as usize];
@@ -78,7 +80,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             &sector_codec,
             plotted_sector.as_mut_slice(),
             io::sink(),
-            piece_receiver_batch_size,
+            &piece_receiver_semaphore,
         ))
         .unwrap();
 

crates/subspace-farmer-components/benches/auditing.rs

@@ -16,6 +16,7 @@ use subspace_core_primitives::{
 };
 use subspace_farmer_components::plotting::plot_sector;
 use subspace_farmer_components::FarmerProtocolInfo;
+use tokio::sync::Semaphore;
 use utils::BenchPieceReceiver;
 
 mod utils;
@@ -54,6 +55,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     };
     let piece_receiver = BenchPieceReceiver::new(piece);
     let piece_receiver_batch_size = 20usize;
+    let piece_receiver_semaphore = Semaphore::new(piece_receiver_batch_size);
 
     let mut group = c.benchmark_group("sector-plotting");
     group.throughput(Throughput::Bytes(PLOT_SECTOR_SIZE));
@@ -69,12 +71,12 @@ fn criterion_benchmark(c: &mut Criterion) {
                 black_box(&sector_codec),
                 black_box(io::sink()),
                 black_box(io::sink()),
-                black_box(piece_receiver_batch_size),
+                black_box(&piece_receiver_semaphore),
             ))
             .unwrap();
         })
     });
-    let piece_receiver_batch_size = 20usize;
+
     let thread_count = current_num_threads() as u64;
     group.throughput(Throughput::Bytes(PLOT_SECTOR_SIZE * thread_count));
     group.bench_function("no-writes-multi-thread", |b| {
@@ -93,7 +95,7 @@ fn criterion_benchmark(c: &mut Criterion) {
                         black_box(&sector_codec),
                         black_box(io::sink()),
                         black_box(io::sink()),
-                        black_box(piece_receiver_batch_size),
+                        black_box(&piece_receiver_semaphore),
                     ))
                     .unwrap();
                 });

crates/subspace-farmer-components/benches/plotting.rs

@@ -20,6 +20,7 @@ use subspace_farmer_components::farming::audit_sector;
 use subspace_farmer_components::file_ext::FileExt;
 use subspace_farmer_components::plotting::plot_sector;
 use subspace_farmer_components::{FarmerProtocolInfo, SectorMetadata};
+use tokio::sync::Semaphore;
 use utils::BenchPieceReceiver;
 
 mod utils;
@@ -66,7 +67,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     let global_challenge = Blake2b256Hash::default();
     let solution_range = SolutionRange::MAX;
     let reward_address = PublicKey::default();
-    let piece_receiver_batch_size = 20usize;
+    let piece_receiver_semaphore = Semaphore::new(usize::MAX);
 
     let (plotted_sector, sector_metadata) = {
         let mut plotted_sector = vec![0u8; PLOT_SECTOR_SIZE as usize];
@@ -82,7 +83,7 @@ pub fn criterion_benchmark(c: &mut Criterion) {
             &sector_codec,
             plotted_sector.as_mut_slice(),
             sector_metadata.as_mut_slice(),
-            piece_receiver_batch_size,
+            &piece_receiver_semaphore,
         ))
         .unwrap();
 

crates/subspace-farmer-components/benches/proving.rs

@@ -6,7 +6,6 @@ use parity_scale_codec::Encode;
 use std::error::Error;
 use std::io;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::Arc;
 use subspace_core_primitives::crypto::kzg;
 use subspace_core_primitives::crypto::kzg::{Commitment, Kzg};
 use subspace_core_primitives::sector_codec::{SectorCodec, SectorCodecError};
@@ -85,7 +84,7 @@ pub async fn plot_sector<PR, S, SM>(
     sector_codec: &SectorCodec,
     mut sector_output: S,
     mut sector_metadata_output: SM,
-    piece_receiver_batch_size: usize,
+    semaphore: &Semaphore,
 ) -> Result<PlottedSector, PlottingError>
 where
     PR: PieceReceiver,
@@ -121,7 +120,7 @@ where
         piece_receiver,
         &piece_indexes,
         cancelled,
-        piece_receiver_batch_size,
+        semaphore,
     )
     .await?;
 
@@ -184,10 +183,8 @@ async fn plot_pieces_in_batches_non_blocking<PR: PieceReceiver>(
     piece_receiver: &PR,
     piece_indexes: &[PieceIndex],
     cancelled: &AtomicBool,
-    piece_receiver_batch_size: usize,
+    semaphore: &Semaphore,
 ) -> Result<(), PlottingError> {
-    let semaphore = Arc::new(Semaphore::new(piece_receiver_batch_size));
-
     let mut pieces_receiving_futures = piece_indexes
         .iter()
         .map(|piece_index| {

crates/subspace-farmer-components/src/plotting.rs

@@ -70,6 +70,15 @@ pub(crate) async fn farm_multi_disk(
 
     let readers_and_pieces = Arc::new(Mutex::new(None));
 
+    info!("Connecting to node RPC at {}", node_rpc_url);
+    let rpc_client = NodeRpcClient::new(&node_rpc_url).await?;
+    let piece_publisher_semaphore = Arc::new(tokio::sync::Semaphore::new(
+        farming_args.piece_receiver_batch_size,
+    ));
+    let piece_receiver_semaphore = Arc::new(tokio::sync::Semaphore::new(
+        farming_args.piece_publisher_batch_size,
+    ));
+
     let (node, mut node_runner) = {
         // TODO: Temporary networking identity derivation from the first disk farm identity.
         let directory = disk_farms
@@ -83,9 +92,6 @@ pub(crate) async fn farm_multi_disk(
 
         if dsn.bootstrap_nodes.is_empty() {
             dsn.bootstrap_nodes = {
-                info!("Connecting to node RPC at {}", node_rpc_url);
-                let rpc_client = NodeRpcClient::new(&node_rpc_url).await?;
-
                 rpc_client
                     .farmer_app_info()
                     .await
@@ -119,8 +125,8 @@ pub(crate) async fn farm_multi_disk(
             rpc_client,
             reward_address,
             dsn_node: node.clone(),
-            piece_receiver_batch_size: farming_args.piece_receiver_batch_size,
-            piece_publisher_batch_size: farming_args.piece_publisher_batch_size,
+            piece_receiver_semaphore: Arc::clone(&piece_receiver_semaphore),
+            piece_publisher_semaphore: Arc::clone(&piece_publisher_semaphore),
         })?;
 
         single_disk_plots.push(single_disk_plot);

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

@@ -271,10 +271,10 @@ pub struct SingleDiskPlotOptions<RC> {
     pub reward_address: PublicKey,
     /// Optional DSN Node.
     pub dsn_node: Node,
-    /// Defines size for the pieces batch of the piece receiving process.
-    pub piece_receiver_batch_size: usize,
-    /// Defines size for the pieces batch of the piece publishing process.
-    pub piece_publisher_batch_size: usize,
+    /// Semaphore to limit concurrency of piece receiving process.
+    pub piece_receiver_semaphore: Arc<tokio::sync::Semaphore>,
+    /// Semaphore to limit concurrency of piece publishing process.
+    pub piece_publisher_semaphore: Arc<tokio::sync::Semaphore>,
 }
 
 /// Errors happening when trying to create/open single disk plot
@@ -468,8 +468,8 @@ impl SingleDiskPlot {
             rpc_client,
             reward_address,
             dsn_node,
-            piece_receiver_batch_size,
-            piece_publisher_batch_size,
+            piece_publisher_semaphore,
+            piece_receiver_semaphore,
         } = options;
 
         // TODO: Account for plot overhead
@@ -648,7 +648,7 @@ impl SingleDiskPlot {
                 let rpc_client = rpc_client.clone();
                 let error_sender = Arc::clone(&error_sender);
                 let piece_publisher =
-                    PieceSectorPublisher::new(dsn_node.clone(), shutting_down.clone());
+                    PieceSectorPublisher::new(dsn_node.clone(), shutting_down.clone(), piece_publisher_semaphore);
 
                 move || {
                     let _tokio_handle_guard = handle.enter();
@@ -709,7 +709,7 @@ impl SingleDiskPlot {
                                 &sector_codec,
                                 sector,
                                 sector_metadata,
-                                piece_receiver_batch_size,
+                                &piece_receiver_semaphore,
                             )) {
                                 Ok(plotted_sector) => {
                                     debug!(%sector_index, "Sector plotted");
@@ -738,7 +738,7 @@ impl SingleDiskPlot {
 
                                 async move {
                                     if let Err(error) = piece_publisher
-                                        .publish_pieces(plotted_sector.piece_indexes, piece_publisher_batch_size)
+                                        .publish_pieces(plotted_sector.piece_indexes)
                                         .await
                                     {
                                         warn!(%sector_index, %error, "Failed to publish pieces to DSN");

crates/subspace-farmer/src/single_disk_plot.rs

@@ -5,8 +5,6 @@ use futures::StreamExt;
 use parity_scale_codec::Encode;
 use std::collections::BTreeSet;
 use std::error::Error;
-use std::future::Future;
-use std::pin::Pin;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::Duration;
@@ -30,13 +28,19 @@ const PUT_PIECE_MAX_INTERVAL: Duration = Duration::from_secs(5);
 pub(crate) struct PieceSectorPublisher {
     dsn_node: Node,
     cancelled: Arc<AtomicBool>,
+    semaphore: Arc<Semaphore>,
 }
 
 impl PieceSectorPublisher {
-    pub(crate) fn new(dsn_node: Node, cancelled: Arc<AtomicBool>) -> Self {
+    pub(crate) fn new(
+        dsn_node: Node,
+        cancelled: Arc<AtomicBool>,
+        semaphore: Arc<Semaphore>,
+    ) -> Self {
         Self {
             dsn_node,
             cancelled,
+            semaphore,
         }
     }
 
@@ -54,15 +58,13 @@ impl PieceSectorPublisher {
     pub(crate) async fn publish_pieces(
         &self,
         piece_indexes: Vec<PieceIndex>,
-        piece_publisher_batch_size: usize,
     ) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
-        let semaphore = Arc::new(Semaphore::new(piece_publisher_batch_size));
-
         let mut pieces_receiving_futures = piece_indexes
             .iter()
             .map(|piece_index| {
                 Box::pin(async {
-                    let _permit = semaphore
+                    let _permit = self
+                        .semaphore
                         .acquire()
                         .await
                         .expect("Should be valid on non-closed semaphore");
@@ -97,12 +99,8 @@ impl PieceSectorPublisher {
             self.check_cancellation()
                 .map_err(backoff::Error::Permanent)?;
 
-            let publish_timeout_result: Result<Result<(), _>, Elapsed> = timeout(
-                PUT_PIECE_TIMEOUT,
-                Box::pin(self.publish_single_piece(piece_index))
-                    as Pin<Box<dyn Future<Output = _> + Send>>,
-            )
-            .await;
+            let publish_timeout_result: Result<Result<(), _>, Elapsed> =
+                timeout(PUT_PIECE_TIMEOUT, self.publish_single_piece(piece_index)).await;
 
             if let Ok(publish_result) = publish_timeout_result {
                 if publish_result.is_ok() {