protocols/kad/query/peers/closest: Make at_capacity use max

protocols/kad/src/query/peers/closest.rs

@@ -369,7 +369,9 @@ impl ClosestPeersIter {
     /// k closest nodes it has not already queried".
     fn at_capacity(&self) -> bool {
         match self.state {
-            State::Stalled => self.num_waiting >= self.config.num_results,
+            State::Stalled => self.num_waiting >= usize::max(
+                self.config.num_results, self.config.parallelism
+            ),
             State::Iterating { .. } => self.num_waiting >= self.config.parallelism,
             State::Finished => true
         }
@@ -691,4 +693,77 @@ mod tests {
 
         QuickCheck::new().tests(10).quickcheck(prop as fn(_) -> _)
     }
+
+    #[test]
+    fn stalled_iter_at_capacity_use_max_of_parallelism_and_num_results() {
+        let target: KeyBytes = Key::from(PeerId::random()).into();
+
+        let mut iter = ClosestPeersIter::with_config(
+            ClosestPeersIterConfig {
+                num_results: 4,
+                parallelism: 8,
+                ..ClosestPeersIterConfig::default()
+            },
+            target,
+            std::iter::empty(),
+        );
+        iter.state = State::Stalled;
+        iter.num_waiting = 7; // Smaller than `parallelism`, but larger than `num_results`.
+
+        assert!(!iter.at_capacity());
+    }
+
+    /// When not making progress for `parallelism` time a [`ClosestPeersIter`] becomes
+    /// [`State::Stalled`]. When stalled an iterator is allowed to make more parallel requests up to
+    /// `num_results`. If `num_results` is smaller than `parallelism` make sure to still allow up to
+    /// `parallelism` requests in-flight.
+    #[test]
+    fn stalled_iter_allows_equal_or_more_than_parallelism_in_flight() {
+        let now = Instant::now();
+
+        let mut num_requests_in_flight = 0;
+
+        let parallelism = 8;
+        let num_results = parallelism / 2;
+        let target: KeyBytes = Key::from(PeerId::random()).into();
+
+        let mut closest_peers = random_peers(100).map(Key::from).collect::<Vec<_>>();
+        closest_peers.sort_unstable_by(|a, b| {
+            target.distance(a).cmp(&target.distance(b))
+        });
+        let mut pool = closest_peers.split_off(K_VALUE.get());
+
+        let mut iter = ClosestPeersIter::with_config(
+            ClosestPeersIterConfig {
+                num_results,
+                parallelism,
+                ..ClosestPeersIterConfig::default()
+            },
+            target,
+            closest_peers.into_iter(),
+        );
+
+        // Have the request for the closest known peer be in-flight for the remainder of the test.
+        // Thereby the iterator never finishes (it ignores timeouts).
+        iter.next(now);
+        num_requests_in_flight += 1;
+
+        while matches!(iter.state, State::Iterating{ .. }) {
+            if let PeersIterState::Waiting(Some(peer)) = iter.next(now) {
+                let peer = peer.clone().into_owned();
+                iter.on_success(
+                    &peer,
+                    pool.pop().map(|p| vec![p.preimage().clone()]).unwrap().into_iter(),
+                );
+            } else {
+                panic!("Expected iterator to yield another peer.");
+            }
+        }
+
+        while matches!(iter.next(now), PeersIterState::Waiting(Some(_))) {
+            num_requests_in_flight += 1;
+        }
+
+        assert_eq!(usize::max(parallelism, num_results), num_requests_in_flight);
+    }
 }