From af9cbb4cba9dfaf7211d7692950bab1a058f08a6 Mon Sep 17 00:00:00 2001
From: Nazar Mokrynskyi <nazar@mokrynskyi.com>
Date: Wed, 23 Oct 2024 17:27:28 +0300
Subject: [PATCH] feat(kad): add `Behavior::find_closest_local_peers()`

---
 protocols/kad/CHANGELOG.md     |  5 ++--
 protocols/kad/src/behaviour.rs | 48 ++++++++++++++++++++--------------
 protocols/kad/src/lib.rs       |  2 +-
 3 files changed, 32 insertions(+), 23 deletions(-)

diff --git a/protocols/kad/CHANGELOG.md b/protocols/kad/CHANGELOG.md
index d0ab7986aad..55d269bf98f 100644
--- a/protocols/kad/CHANGELOG.md
+++ b/protocols/kad/CHANGELOG.md
@@ -1,10 +1,11 @@
 ## 0.47.0
 
-- Expose a kad query facility allowing specify num_results dynamicly.
+- Expose a kad query facility allowing specify num_results dynamicaly.
   See [PR 5555](https://github.com/libp2p/rust-libp2p/pull/5555).
 - Add `mode` getter on `Behaviour`.
   See [PR 5573](https://github.com/libp2p/rust-libp2p/pull/5573).
-  
+- Add `Behavior::find_closest_local_peers()`.
+  See [PR 5645](https://github.com/libp2p/rust-libp2p/pull/5645).
 
 ## 0.46.2
 
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
index 0b15e507ba4..84133d31acb 100644
--- a/protocols/kad/src/behaviour.rs
+++ b/protocols/kad/src/behaviour.rs
@@ -771,7 +771,8 @@ where
         self.queries.add_iter_closest(target, peer_keys, info)
     }
 
-    /// Returns closest peers to the given key; takes peers from local routing table only.
+    /// Returns all peers ordered by distance to the given key; takes peers from local routing table
+    /// only.
     pub fn get_closest_local_peers<'a, K: Clone>(
         &'a mut self,
         key: &'a kbucket::Key<K>,
@@ -779,6 +780,23 @@ where
         self.kbuckets.closest_keys(key)
     }
 
+    /// Finds the closest peers to a `key` in the context of a request by the `source` peer, such
+    /// that the `source` peer is never included in the result.
+    ///
+    /// Takes peers from local routing table only. Only returns number of peers equal to configured
+    /// replication factor.
+    pub fn find_closest_local_peers<'a, K: Clone>(
+        &'a mut self,
+        key: &'a kbucket::Key<K>,
+        source: &'a PeerId,
+    ) -> impl Iterator<Item = KadPeer> + 'a {
+        self.kbuckets
+            .closest(key)
+            .filter(move |e| e.node.key.preimage() != source)
+            .take(self.queries.config().replication_factor.get())
+            .map(KadPeer::from)
+    }
+
     /// Performs a lookup for a record in the DHT.
     ///
     /// The result of this operation is delivered in a
@@ -1212,22 +1230,6 @@ where
         }
     }
 
-    /// Finds the closest peers to a `target` in the context of a request by
-    /// the `source` peer, such that the `source` peer is never included in the
-    /// result.
-    fn find_closest<T: Clone>(
-        &mut self,
-        target: &kbucket::Key<T>,
-        source: &PeerId,
-    ) -> Vec<KadPeer> {
-        self.kbuckets
-            .closest(target)
-            .filter(|e| e.node.key.preimage() != source)
-            .take(self.queries.config().replication_factor.get())
-            .map(KadPeer::from)
-            .collect()
-    }
-
     /// Collects all peers who are known to be providers of the value for a given `Multihash`.
     fn provider_peers(&mut self, key: &record::Key, source: &PeerId) -> Vec<KadPeer> {
         let kbuckets = &mut self.kbuckets;
@@ -2300,7 +2302,9 @@ where
             }
 
             HandlerEvent::FindNodeReq { key, request_id } => {
-                let closer_peers = self.find_closest(&kbucket::Key::new(key), &source);
+                let closer_peers = self
+                    .find_closest_local_peers(&kbucket::Key::new(key), &source)
+                    .collect::<Vec<_>>();
 
                 self.queued_events
                     .push_back(ToSwarm::GenerateEvent(Event::InboundRequest {
@@ -2328,7 +2332,9 @@ where
 
             HandlerEvent::GetProvidersReq { key, request_id } => {
                 let provider_peers = self.provider_peers(&key, &source);
-                let closer_peers = self.find_closest(&kbucket::Key::new(key), &source);
+                let closer_peers = self
+                    .find_closest_local_peers(&kbucket::Key::new(key), &source)
+                    .collect::<Vec<_>>();
 
                 self.queued_events
                     .push_back(ToSwarm::GenerateEvent(Event::InboundRequest {
@@ -2422,7 +2428,9 @@ where
                     None => None,
                 };
 
-                let closer_peers = self.find_closest(&kbucket::Key::new(key), &source);
+                let closer_peers = self
+                    .find_closest_local_peers(&kbucket::Key::new(key), &source)
+                    .collect::<Vec<_>>();
 
                 self.queued_events
                     .push_back(ToSwarm::GenerateEvent(Event::InboundRequest {
diff --git a/protocols/kad/src/lib.rs b/protocols/kad/src/lib.rs
index 681d135f79b..060bfc518e4 100644
--- a/protocols/kad/src/lib.rs
+++ b/protocols/kad/src/lib.rs
@@ -69,7 +69,7 @@ pub use behaviour::{
 pub use kbucket::{
     Distance as KBucketDistance, EntryView, KBucketRef, Key as KBucketKey, NodeStatus,
 };
-pub use protocol::ConnectionType;
+pub use protocol::{ConnectionType, KadPeer};
 pub use query::QueryId;
 pub use record::{store, Key as RecordKey, ProviderRecord, Record};