Merge pull request #715 from subspace/move-get-pieces-by-range

Move get pieces by range

Author: i1i1

crates/subspace-farmer/src/dsn.rs

@@ -1,16 +1,35 @@
-use futures::{Stream, StreamExt};
+use futures::{SinkExt, Stream, StreamExt};
 use num_traits::{WrappingAdd, WrappingSub};
 use std::ops::Range;
 use subspace_core_primitives::{
     FlatPieces, PieceIndex, PieceIndexHash, PublicKey, PIECE_SIZE, U256,
 };
-use subspace_networking::PiecesToPlot;
+use subspace_networking::libp2p::core::multihash::{Code, MultihashDigest};
+use subspace_networking::{PiecesByRangeRequest, PiecesByRangeResponse, PiecesToPlot};
+use tracing::{debug, trace, warn};
 
+#[cfg(test)]
+mod pieces_by_range_tests;
 #[cfg(test)]
 mod tests;
 
+const PIECES_CHANNEL_BUFFER_SIZE: usize = 20;
+
 pub type PieceIndexHashNumber = U256;
 
+#[derive(Debug, thiserror::Error)]
+pub enum GetPiecesByRangeError {
+    /// Cannot find closest pieces by range.
+    #[error("Cannot find closest pieces by range")]
+    NoClosestPiecesFound,
+    /// Cannot get closest peers from DHT.
+    #[error("Cannot get closest peers from DHT")]
+    CannotGetClosestPeers(#[source] subspace_networking::GetClosestPeersError),
+    /// Node runner was dropped, impossible to get pieces by range.
+    #[error("Node runner was dropped, impossible to get pieces by range")]
+    NodeRunnerDropped(#[source] subspace_networking::SendRequestError),
+}
+
 /// Options for syncing
 pub struct SyncOptions {
     /// Max plot size from node (in bytes)
@@ -55,13 +74,91 @@ impl DSNSync for NoSync {
 #[async_trait::async_trait]
 impl DSNSync for subspace_networking::Node {
     type Stream = futures::channel::mpsc::Receiver<PiecesToPlot>;
-    type Error = subspace_networking::GetPiecesByRangeError;
+    type Error = GetPiecesByRangeError;
 
     async fn get_pieces(
         &mut self,
         Range { start, end }: Range<PieceIndexHash>,
     ) -> Result<Self::Stream, Self::Error> {
-        self.get_pieces_by_range(start, end).await
+        // calculate the middle of the range (big endian)
+        let middle = {
+            let start = PieceIndexHashNumber::from(start);
+            let end = PieceIndexHashNumber::from(end);
+            // min + (max - min) / 2
+            (start / 2 + end / 2).to_be_bytes()
+        };
+
+        // obtain closest peers to the middle of the range
+        let key = Code::Identity.digest(&middle);
+        let peers = self
+            .get_closest_peers(key)
+            .await
+            .map_err(GetPiecesByRangeError::CannotGetClosestPeers)?;
+
+        // select first peer for the piece-by-range protocol
+        let peer_id = *peers
+            .first()
+            .ok_or(GetPiecesByRangeError::NoClosestPiecesFound)?;
+
+        trace!(%peer_id, ?start, ?end, "Peer found");
+
+        // prepare stream channel
+        let (mut tx, rx) =
+            futures::channel::mpsc::channel::<PiecesToPlot>(PIECES_CHANNEL_BUFFER_SIZE);
+
+        // populate resulting stream in a separate async task
+        let node = self.clone();
+        tokio::spawn(async move {
+            // indicates the next starting point for a request
+            let mut starting_index_hash = start;
+            loop {
+                trace!(
+                    "Sending 'Piece-by-range' request to {} with {:?}",
+                    peer_id,
+                    starting_index_hash
+                );
+                // request data by range
+                let response = node
+                    .send_generic_request(
+                        peer_id,
+                        PiecesByRangeRequest {
+                            start: starting_index_hash,
+                            end,
+                        },
+                    )
+                    .await
+                    .map_err(GetPiecesByRangeError::NodeRunnerDropped);
+
+                // send the result to the stream and exit on any error
+                match response {
+                    Ok(PiecesByRangeResponse {
+                        pieces,
+                        next_piece_index_hash,
+                    }) => {
+                        // send last response data stream to the result stream
+                        if !pieces.piece_indexes.is_empty() && tx.send(pieces).await.is_err() {
+                            warn!("Piece-by-range request channel was closed.");
+                            break;
+                        }
+
+                        // prepare the next starting point for data
+                        if let Some(next_piece_index_hash) = next_piece_index_hash {
+                            debug_assert_ne!(starting_index_hash, next_piece_index_hash);
+                            starting_index_hash = next_piece_index_hash;
+                        } else {
+                            // exit loop if the last response showed no remaining data
+                            break;
+                        }
+                    }
+                    Err(err) => {
+                        debug!(%err, "Piece-by-range request failed");
+                        break;
+                    }
+                }
+            }
+        });
+
+        Ok(rx)
     }
 }
 

crates/subspace-networking/src/request_handlers/pieces_by_range/tests.rs renamed to crates/subspace-farmer/src/dsn/pieces_by_range_tests.rs

@@ -1,19 +1,20 @@
-use crate::{
-    Config, PiecesByRangeRequest, PiecesByRangeRequestHandler, PiecesByRangeResponse, PiecesToPlot,
-};
+use super::DSNSync;
 use futures::channel::{mpsc, oneshot};
 use futures::{SinkExt, StreamExt};
-use libp2p::multiaddr::Protocol;
 use parking_lot::Mutex;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use subspace_core_primitives::{crypto, FlatPieces, Piece, PieceIndexHash};
+use subspace_networking::libp2p::multiaddr::Protocol;
+use subspace_networking::{
+    Config, PiecesByRangeRequest, PiecesByRangeRequestHandler, PiecesByRangeResponse, PiecesToPlot,
+};
 
 #[tokio::test]
 async fn pieces_by_range_protocol_smoke() {
     let request = PiecesByRangeRequest {
-        from: PieceIndexHash::from([1u8; 32]),
-        to: PieceIndexHash::from([1u8; 32]),
+        start: PieceIndexHash::from([1u8; 32]),
+        end: PieceIndexHash::from([1u8; 32]),
     };
 
     let piece_bytes: Vec<u8> = Piece::default().into();
@@ -41,7 +42,7 @@ async fn pieces_by_range_protocol_smoke() {
         })],
         ..Config::with_generated_keypair()
     };
-    let (node_1, mut node_runner_1) = crate::create(config_1).await.unwrap();
+    let (node_1, mut node_runner_1) = subspace_networking::create(config_1).await.unwrap();
 
     let (node_1_address_sender, node_1_address_receiver) = oneshot::channel();
     let on_new_listener_handler = node_1.on_new_listener(Arc::new({
@@ -72,7 +73,7 @@ async fn pieces_by_range_protocol_smoke() {
         ..Config::with_generated_keypair()
     };
 
-    let (node_2, mut node_runner_2) = crate::create(config_2).await.unwrap();
+    let (node_2, mut node_runner_2) = subspace_networking::create(config_2).await.unwrap();
     tokio::spawn(async move {
         node_runner_2.run().await;
     });
@@ -130,7 +131,7 @@ async fn get_pieces_by_range_smoke() {
                 })
             } else {
                 // New request starts from from the previous response.
-                assert_eq!(req.from, piece_index_continue);
+                assert_eq!(req.start, piece_index_continue);
 
                 Some(PiecesByRangeResponse {
                     pieces: response_data[request_index].clone(),
@@ -140,7 +141,7 @@ async fn get_pieces_by_range_smoke() {
         })],
         ..Config::with_generated_keypair()
     };
-    let (node_1, mut node_runner_1) = crate::create(config_1).await.unwrap();
+    let (node_1, mut node_runner_1) = subspace_networking::create(config_1).await.unwrap();
 
     let (node_1_address_sender, node_1_address_receiver) = oneshot::channel();
     let on_new_listener_handler = node_1.on_new_listener(Arc::new({
@@ -171,13 +172,13 @@ async fn get_pieces_by_range_smoke() {
         ..Config::with_generated_keypair()
     };
 
-    let (node_2, mut node_runner_2) = crate::create(config_2).await.unwrap();
+    let (mut node_2, mut node_runner_2) = subspace_networking::create(config_2).await.unwrap();
     tokio::spawn(async move {
         node_runner_2.run().await;
     });
 
     let mut stream = node_2
-        .get_pieces_by_range(piece_index_from, piece_index_end)
+        .get_pieces(piece_index_from..piece_index_end)
         .await
         .unwrap();
 

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

@@ -1,4 +1,4 @@
-use super::{sync, DSNSync, NoSync, PieceIndexHashNumber, SyncOptions};
+use super::{sync, DSNSync, NoSync, PieceIndexHashNumber, PiecesToPlot, SyncOptions};
 use crate::bench_rpc_client::{BenchRpcClient, BENCH_FARMER_PROTOCOL_INFO};
 use crate::legacy_multi_plots_farm::{LegacyMultiPlotsFarm, Options as MultiFarmingOptions};
 use crate::single_plot_farm::PlotFactoryOptions;
@@ -18,7 +18,6 @@ use subspace_core_primitives::{
     PieceIndex, PieceIndexHash, RootBlock, Sha256Hash, PIECE_SIZE, U256,
 };
 use subspace_networking::libp2p::multiaddr::Protocol;
-use subspace_networking::PiecesToPlot;
 use tempfile::TempDir;
 
 #[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]

crates/subspace-farmer/src/single_plot_farm.rs

@@ -3,8 +3,7 @@ pub mod dsn_archiving;
 mod tests;
 
 use crate::commitments::{CommitmentError, Commitments};
-use crate::dsn;
-use crate::dsn::{PieceIndexHashNumber, SyncOptions};
+use crate::dsn::{self, PieceIndexHashNumber, SyncOptions};
 use crate::farming::Farming;
 use crate::identity::Identity;
 use crate::object_mappings::ObjectMappings;
@@ -451,13 +450,14 @@ impl SinglePlotFarm {
                 let codec = codec.clone();
 
                 // TODO: also ask for how many pieces to read
-                move |&PiecesByRangeRequest { from, to }| {
-                    let mut pieces_and_indexes =
-                        plot.get_sequential_pieces(from, SYNC_PIECES_AT_ONCE).ok()?;
+                move |&PiecesByRangeRequest { start, end }| {
+                    let mut pieces_and_indexes = plot
+                        .get_sequential_pieces(start, SYNC_PIECES_AT_ONCE)
+                        .ok()?;
 
                     let next_piece_index_hash = if let Some(idx) = pieces_and_indexes
                         .iter()
-                        .position(|(piece_index, _)| PieceIndexHash::from_index(*piece_index) > to)
+                        .position(|(piece_index, _)| PieceIndexHash::from_index(*piece_index) > end)
                     {
                         pieces_and_indexes.truncate(idx);
                         None

crates/subspace-networking/examples/get-pieces-complex.rs renamed to crates/subspace-networking/examples/get-peers-complex.rs

@@ -1,14 +1,10 @@
 use futures::channel::oneshot;
-use futures::StreamExt;
 use libp2p::multiaddr::Protocol;
-use libp2p::{identity, PeerId};
+use libp2p::PeerId;
 use parking_lot::Mutex;
 use std::sync::Arc;
 use std::time::Duration;
-use subspace_core_primitives::{FlatPieces, Piece, PieceIndexHash};
-use subspace_networking::{
-    Config, PiecesByRangeRequestHandler, PiecesByRangeResponse, PiecesToPlot,
-};
+use subspace_networking::Config;
 
 #[tokio::main]
 async fn main() {
@@ -20,36 +16,11 @@ async fn main() {
     const TOTAL_NODE_COUNT: usize = 100;
     const EXPECTED_NODE_INDEX: usize = 75;
 
-    let expected_response = {
-        let piece_bytes: Vec<u8> = Piece::default().into();
-        let flat_pieces = FlatPieces::try_from(piece_bytes).unwrap();
-        let pieces = PiecesToPlot {
-            piece_indexes: vec![1],
-            pieces: flat_pieces,
-        };
-
-        PiecesByRangeResponse {
-            pieces,
-            next_piece_index_hash: None,
-        }
-    };
-
     for i in 0..TOTAL_NODE_COUNT {
-        let local_response = expected_response.clone();
         let config = Config {
             bootstrap_nodes: bootstrap_nodes.clone(),
             listen_on: vec!["/ip4/0.0.0.0/tcp/0".parse().unwrap()],
             allow_non_globals_in_dht: true,
-            request_response_protocols: vec![PiecesByRangeRequestHandler::create(move |_| {
-                if i != EXPECTED_NODE_INDEX {
-                    return None;
-                }
-
-                println!("Sending response from Node Index {}... ", i);
-
-                std::thread::sleep(Duration::from_secs(1));
-                Some(local_response.clone())
-            })],
             ..Config::with_generated_keypair()
         };
         let (node, mut node_runner) = subspace_networking::create(config).await.unwrap();
@@ -91,7 +62,6 @@ async fn main() {
         bootstrap_nodes,
         listen_on: vec!["/ip4/0.0.0.0/tcp/0".parse().unwrap()],
         allow_non_globals_in_dht: true,
-        request_response_protocols: vec![PiecesByRangeRequestHandler::create(|_request| None)],
         ..Config::with_generated_keypair()
     };
 
@@ -106,37 +76,11 @@ async fn main() {
 
     tokio::time::sleep(Duration::from_secs(1)).await;
 
-    let encoding = expected_node_id.as_ref().digest();
-    let public_key = identity::PublicKey::from_protobuf_encoding(encoding)
-        .expect("Invalid public key from PeerId.");
-    let peer_id_public_key = if let identity::PublicKey::Sr25519(pk) = public_key {
-        pk.encode()
-    } else {
-        panic!("Expected PublicKey::Sr25519")
-    };
+    let peer_id = node
+        .get_closest_peers(expected_node_id.into())
+        .await
+        .unwrap()[0];
+    assert_eq!(peer_id, expected_node_id);
 
-    // create a range from expected peer's public key
-    let from = {
-        let mut buf = peer_id_public_key;
-        buf[16] = 0;
-        PieceIndexHash::from(buf)
-    };
-    let to = {
-        let mut buf = peer_id_public_key;
-        buf[16] = 50;
-        PieceIndexHash::from(buf)
-    };
-
-    let stream_future = node.get_pieces_by_range(from, to);
-    let mut stream = stream_future.await.unwrap();
-    if let Some(value) = stream.next().await {
-        if value != expected_response.pieces {
-            panic!("UNEXPECTED RESPONSE")
-        }
-
-        println!("Received expected response.");
-    }
-
-    tokio::time::sleep(Duration::from_secs(1)).await;
     println!("Exiting..");
 }