refactor to use Index structure and emit multiple providers at once

libp2p · Jul 26, 2022 · dc5db66 · dc5db66
1 parent b224c59
commit dc5db66
Show file tree

Hide file tree

Showing 7 changed files with 226 additions and 173 deletions.
diff --git a/examples/distributed-key-value-store.rs b/examples/distributed-key-value-store.rs
@@ -53,7 +53,7 @@ use libp2p::{
     swarm::{NetworkBehaviourEventProcess, SwarmEvent},
     NetworkBehaviour, PeerId, Swarm,
 };
-use libp2p_kad::{GetProvidersProgress, QueryProgress};
+use libp2p_kad::GetProvidersOk;
 use std::error::Error;
 
 #[async_std::main]
@@ -91,15 +91,15 @@ async fn main() -> Result<(), Box<dyn Error>> {
         fn inject_event(&mut self, message: KademliaEvent) {
             match message {
                 KademliaEvent::OutboundQueryProgressed { result, .. } => match result {
-                    QueryProgress::GetProviders(GetProvidersProgress { key, provider, .. }) => {
-                        println!(
-                            "Peer {:?} provides key {:?}",
-                            provider,
-                            std::str::from_utf8(key.as_ref()).unwrap()
-                        );
+                    QueryResult::GetProviders(Ok(GetProvidersOk { key, providers, .. })) => {
+                        for provider in &providers {
+                            println!(
+                                "Peer {:?} provides key {:?}",
+                                provider,
+                                std::str::from_utf8(key.as_ref()).unwrap()
+                            );
+                        }
                     }
-                },
-                KademliaEvent::OutboundQueryCompleted { result, .. } => match result {
                     QueryResult::GetProviders(Err(err)) => {
                         eprintln!("Failed to get providers: {:?}", err);
                     }

diff --git a/examples/file-sharing.rs b/examples/file-sharing.rs
@@ -219,7 +219,7 @@ mod network {
     };
     use libp2p::swarm::{ConnectionHandlerUpgrErr, SwarmBuilder, SwarmEvent};
     use libp2p::{NetworkBehaviour, Swarm};
-    use libp2p_kad::{GetProvidersProgress, QueryProgress};
+    use libp2p_kad::GetProvidersOk;
     use std::collections::{HashMap, HashSet};
     use std::iter;
 
@@ -325,12 +325,16 @@ mod network {
 
         /// Find the providers for the given file on the DHT.
         pub async fn get_providers(&mut self, file_name: String) -> HashSet<PeerId> {
-            let (sender, receiver) = mpsc::channel(0);
+            let (sender, mut receiver) = mpsc::channel(0);
             self.sender
                 .send(Command::GetProviders { file_name, sender })
                 .await
                 .expect("Command receiver not to be dropped.");
-            receiver.collect().await
+            let mut out = HashSet::new();
+            while let Some(h) = receiver.next().await {
+                out.extend(h);
+            }
+            out
         }
 
         /// Request the content of the given file from the given peer.
@@ -366,7 +370,7 @@ mod network {
         event_sender: mpsc::Sender<Event>,
         pending_dial: HashMap<PeerId, oneshot::Sender<Result<(), Box<dyn Error + Send>>>>,
         pending_start_providing: HashMap<QueryId, oneshot::Sender<()>>,
-        pending_get_providers: HashMap<QueryId, mpsc::Sender<PeerId>>,
+        pending_get_providers: HashMap<QueryId, mpsc::Sender<HashSet<PeerId>>>,
         pending_request_file:
             HashMap<RequestId, oneshot::Sender<Result<String, Box<dyn Error + Send>>>>,
     }
@@ -410,7 +414,7 @@ mod network {
         ) {
             match event {
                 SwarmEvent::Behaviour(ComposedEvent::Kademlia(
-                    KademliaEvent::OutboundQueryCompleted {
+                    KademliaEvent::OutboundQueryProgressed {
                         id,
                         result: QueryResult::StartProviding(_),
                         ..
@@ -425,18 +429,18 @@ mod network {
                 SwarmEvent::Behaviour(ComposedEvent::Kademlia(
                     KademliaEvent::OutboundQueryProgressed {
                         id,
-                        result: QueryProgress::GetProviders(GetProvidersProgress { provider, .. }),
+                        result: QueryResult::GetProviders(Ok(GetProvidersOk { providers, .. })),
                         ..
                     },
                 )) => {
                     let _ = self
                         .pending_get_providers
                         .get_mut(&id)
                         .expect("Completed query to be previously pending.")
-                        .send(provider);
+                        .send(providers);
                 }
                 SwarmEvent::Behaviour(ComposedEvent::Kademlia(
-                    KademliaEvent::OutboundQueryCompleted {
+                    KademliaEvent::OutboundQueryProgressed {
                         id,
                         result: QueryResult::GetProviders(..),
                         ..
@@ -634,7 +638,7 @@ mod network {
         },
         GetProviders {
             file_name: String,
-            sender: mpsc::Sender<PeerId>,
+            sender: mpsc::Sender<HashSet<PeerId>>,
         },
         RequestFile {
             file_name: String,

diff --git a/examples/ipfs-kad.rs b/examples/ipfs-kad.rs
@@ -85,7 +85,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
     task::block_on(async move {
         loop {
             let event = swarm.select_next_some().await;
-            if let SwarmEvent::Behaviour(KademliaEvent::OutboundQueryCompleted {
+            if let SwarmEvent::Behaviour(KademliaEvent::OutboundQueryProgressed {
                 result: QueryResult::GetClosestPeers(result),
                 ..
             }) = event

diff --git a/misc/metrics/src/kad.rs b/misc/metrics/src/kad.rs
@@ -18,6 +18,7 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+use libp2p_kad::GetProvidersOk;
 use prometheus_client::encoding::text::Encode;
 use prometheus_client::metrics::counter::Counter;
 use prometheus_client::metrics::family::Family;
@@ -162,7 +163,7 @@ impl Metrics {
 impl super::Recorder<libp2p_kad::KademliaEvent> for Metrics {
     fn record(&self, event: &libp2p_kad::KademliaEvent) {
         match event {
-            libp2p_kad::KademliaEvent::OutboundQueryCompleted { result, stats, .. } => {
+            libp2p_kad::KademliaEvent::OutboundQueryProgressed { result, stats, .. } => {
                 self.query_result_num_requests
                     .get_or_create(&result.into())
                     .observe(stats.num_requests().into());
@@ -200,7 +201,10 @@ impl super::Recorder<libp2p_kad::KademliaEvent> for Metrics {
                         }
                     },
                     libp2p_kad::QueryResult::GetProviders(result) => match result {
-                        Ok(ok) => {}
+                        Ok(GetProvidersOk { providers, .. }) => {
+                            self.query_result_get_providers_ok
+                                .observe(providers.len() as f64);
+                        }
                         Err(error) => {
                             self.query_result_get_providers_error
                                 .get_or_create(&error.into())
@@ -210,11 +214,6 @@ impl super::Recorder<libp2p_kad::KademliaEvent> for Metrics {
                     _ => {}
                 }
             }
-            libp2p_kad::KademliaEvent::OutboundQueryProgressed { result, .. } => match result {
-                libp2p_kad::QueryProgress::GetProviders(_) => {
-                    self.query_result_get_providers_ok.observe(1.);
-                }
-            },
             libp2p_kad::KademliaEvent::RoutingUpdated {
                 is_new_peer,
                 old_peer,