From 3d16353ef972fe072d167bc198cb81c898e9198d Mon Sep 17 00:00:00 2001
From: maqi <qi.ma@maidsafe.net>
Date: Wed, 3 Apr 2024 15:39:39 +0800
Subject: [PATCH] fix(kad): close progress whenever closer in range

Pull-Request: #4934.
---
 protocols/kad/CHANGELOG.md               |  2 ++
 protocols/kad/src/query/peers/closest.rs | 25 ++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/protocols/kad/CHANGELOG.md b/protocols/kad/CHANGELOG.md
index c300b3127da..0a561b14448 100644
--- a/protocols/kad/CHANGELOG.md
+++ b/protocols/kad/CHANGELOG.md
@@ -4,6 +4,8 @@
   See [PR 5270](https://github.com/libp2p/rust-libp2p/pull/5270)
 - Update to DHT republish interval and expiration time defaults to 22h and 48h respectively, rationale in [libp2p/specs#451](https://github.com/libp2p/specs/pull/451)
   See [PR 3230](https://github.com/libp2p/rust-libp2p/pull/3230)
+- QueryClose progress whenever closer in range, instead of having to be the closest.
+  See [PR 4934](https://github.com/libp2p/rust-libp2p/pull/4934).
 
 ## 0.45.4
 
diff --git a/protocols/kad/src/query/peers/closest.rs b/protocols/kad/src/query/peers/closest.rs
index 537da6a9d7d..3ee12b98244 100644
--- a/protocols/kad/src/query/peers/closest.rs
+++ b/protocols/kad/src/query/peers/closest.rs
@@ -174,6 +174,20 @@ impl ClosestPeersIter {
             },
         }
 
+        let mut cur_range = distance;
+        let num_results = self.config.num_results.get();
+        // furthest_peer is the furthest peer in range among the closest_peers
+        let furthest_peer = self
+            .closest_peers
+            .iter()
+            .enumerate()
+            .nth(num_results - 1)
+            .map(|(_, peer)| peer)
+            .or_else(|| self.closest_peers.iter().last());
+        if let Some((dist, _)) = furthest_peer {
+            cur_range = *dist;
+        }
+
         // Incorporate the reported closer peers into the iterator.
         //
         // The iterator makes progress if:
@@ -189,9 +203,16 @@ impl ClosestPeersIter {
                 key,
                 state: PeerState::NotContacted,
             };
-            self.closest_peers.entry(distance).or_insert(peer);
 
-            progress = self.closest_peers.keys().next() == Some(&distance) || progress;
+            let is_first_insert = match self.closest_peers.entry(distance) {
+                Entry::Occupied(_) => false,
+                Entry::Vacant(entry) => {
+                    entry.insert(peer);
+                    true
+                }
+            };
+
+            progress = (is_first_insert && distance < cur_range) || progress;
         }
 
         // Update the iterator state.