chore(kad): expose a kad query facility allowing dynamic num_results

Cargo.toml

@@ -86,7 +86,7 @@ libp2p-floodsub = { version = "0.45.0", path = "protocols/floodsub" }
 libp2p-gossipsub = { version = "0.47.0", path = "protocols/gossipsub" }
 libp2p-identify = { version = "0.45.0", path = "protocols/identify" }
 libp2p-identity = { version = "0.2.9" }
-libp2p-kad = { version = "0.46.0", path = "protocols/kad" }
+libp2p-kad = { version = "0.46.1", path = "protocols/kad" }
 libp2p-mdns = { version = "0.46.0", path = "protocols/mdns" }
 libp2p-memory-connection-limits = { version = "0.2.0", path = "misc/memory-connection-limits" }
 libp2p-metrics = { version = "0.14.2", path = "misc/metrics" }

protocols/kad/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 0.46.1
+
+- Expose a kad query facility allowing specify num_results dynamicly.
+  See [PR 5555](https://github.com/libp2p/rust-libp2p/pull/5555).
+
 ## 0.46.0
 
 - Included multiaddresses of found peers alongside peer IDs in `GetClosestPeers` query results.

protocols/kad/Cargo.toml

@@ -3,7 +3,7 @@ name = "libp2p-kad"
 edition = "2021"
 rust-version = { workspace = true }
 description = "Kademlia protocol for libp2p"
-version = "0.46.0"
+version = "0.46.1"
 authors = ["Parity Technologies <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"

protocols/kad/src/behaviour.rs

@@ -732,6 +732,26 @@ where
     /// The result of the query is delivered in a
     /// [`Event::OutboundQueryProgressed{QueryResult::GetClosestPeers}`].
     pub fn get_closest_peers<K>(&mut self, key: K) -> QueryId
+    where
+        K: Into<kbucket::Key<K>> + Into<Vec<u8>> + Clone,
+    {
+        self.get_closest_peers_inner(key, None)
+    }
+
+    /// Initiates an iterative query for the closest peers to the given key.
+    ///
+    /// The result of the query is delivered in a
+    /// [`Event::OutboundQueryProgressed{QueryResult::GetClosestPeers}`].
+    ///
+    /// The expected responding peers is specified by `num_results`
+    pub fn get_closest_peers_num_results<K>(&mut self, key: K, num_results: NonZeroUsize) -> QueryId
+    where
+        K: Into<kbucket::Key<K>> + Into<Vec<u8>> + Clone,
+    {
+        self.get_closest_peers_inner(key, Some(num_results))
+    }
+
+    fn get_closest_peers_inner<K>(&mut self, key: K, num_results: Option<NonZeroUsize>) -> QueryId
     where
         K: Into<kbucket::Key<K>> + Into<Vec<u8>> + Clone,
     {
@@ -742,7 +762,13 @@ where
             step: ProgressStep::first(),
         };
         let peer_keys: Vec<kbucket::Key<PeerId>> = self.kbuckets.closest_keys(&target).collect();
-        self.queries.add_iter_closest(target, peer_keys, info)
+
+        if let Some(num_results) = num_results {
+            self.queries
+                .add_iter_closest_num_results(target, peer_keys, info, num_results)
+        } else {
+            self.queries.add_iter_closest(target, peer_keys, info)
+        }
     }
 
     /// Returns closest peers to the given key; takes peers from local routing table only.

protocols/kad/src/query.rs

@@ -127,6 +127,24 @@ impl QueryPool {
         id
     }
 
+    /// Adds a query to the pool that iterates towards the closest peers to the target.
+    /// With `num_results` dynamic specified.
+    pub(crate) fn add_iter_closest_num_results<T, I>(
+        &mut self,
+        target: T,
+        peers: I,
+        info: QueryInfo,
+        num_results: NonZeroUsize,
+    ) -> QueryId
+    where
+        T: Into<KeyBytes> + Clone,
+        I: IntoIterator<Item = Key<PeerId>>,
+    {
+        let id = self.next_query_id();
+        self.continue_iter_closest_inner(id, target, peers, info, num_results);
+        id
+    }
+
     /// Adds a query to the pool that iterates towards the closest peers to the target.
     pub(crate) fn continue_iter_closest<T, I>(
         &mut self,
@@ -137,9 +155,23 @@ impl QueryPool {
     ) where
         T: Into<KeyBytes> + Clone,
         I: IntoIterator<Item = Key<PeerId>>,
+    {
+        self.continue_iter_closest_inner(id, target, peers, info, self.config.replication_factor);
+    }
+
+    fn continue_iter_closest_inner<T, I>(
+        &mut self,
+        id: QueryId,
+        target: T,
+        peers: I,
+        info: QueryInfo,
+        num_results: NonZeroUsize,
+    ) where
+        T: Into<KeyBytes> + Clone,
+        I: IntoIterator<Item = Key<PeerId>>,
     {
         let cfg = ClosestPeersIterConfig {
-            num_results: self.config.replication_factor,
+            num_results,
             parallelism: self.config.parallelism,
             ..ClosestPeersIterConfig::default()
         };