diff --git a/protocols/kad/src/query/peers/closest.rs b/protocols/kad/src/query/peers/closest.rs index f4e37e3d12e..987d0c36c85 100644 --- a/protocols/kad/src/query/peers/closest.rs +++ b/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,29 @@ mod tests { QuickCheck::new().tests(10).quickcheck(prop as fn(_) -> _) } + + #[test] + fn stalled_at_capacity() { + fn prop(mut iter: ClosestPeersIter) { + iter.state = State::Stalled; + + for i in 0..usize::max(iter.config.parallelism, iter.config.num_results) { + iter.num_waiting = i; + assert!( + !iter.at_capacity(), + "Iterator should not be at capacity if less than \ + `max(parallelism, num_results)` requests are waiting.", + ) + } + + iter.num_waiting = usize::max(iter.config.parallelism, iter.config.num_results); + assert!( + iter.at_capacity(), + "Iterator should be at capacity if `max(parallelism, num_results)` requests are \ + waiting.", + ) + } + + QuickCheck::new().tests(10).quickcheck(prop as fn(_)) + } }