chore(kad): continuous lookup pipeline

gmelodie · gmelodie · commit ce80d3837b50 · 2026-07-17T15:30:49.000-03:00
diff --git a/libp2p/protocols/kademlia/find.nim b/libp2p/protocols/kademlia/find.nim
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0 OR MIT
 # Copyright (c) Status Research & Development GmbH
 
-import std/[tables, sequtils, algorithm, net]
+import std/[tables, sequtils, algorithm, net, sets]
 import chronos, chronicles, results
 import ../../[peerid, peerinfo, switch, multihash, peeraddrpolicy, wire]
 import ../protocol
@@ -21,7 +21,6 @@ type LookupState* = ref object
   shortlist*: Table[PeerId, XorDistance]
   responded*: Table[PeerId, RespondedStatus]
   attempts*: Table[PeerId, int]
-  inflight*: Table[PeerId, seq[FutureBase]]
 
 type DispatchProc* = proc(
   kad: KadDHT, peer: PeerId, target: Key
@@ -366,77 +365,88 @@ proc findNodeDispatch*(
 ): Future[Result[Message, string]] {.async: (raises: [CancelledError]), gcsafe.} =
   return await dispatchFindNode(kad, peer, target)
 
-proc lookOnce*(
-    kad: KadDHT,
-    state: LookupState,
-    rtable: RoutingTable,
-    dispatch: DispatchProc,
-    onReply: ReplyHandler,
-): Future[bool] {.async: (raises: [CancelledError]).} =
-  ## Run a single round of the iterative lookup against ``state.target``.
-  ## Returns ``false`` when no further peers can be queried, signalling the
-  ## caller to stop driving the lookup.
-  let toQuery = state.selectCloserPeers(kad.config.alpha)
-  if toQuery.len() == 0:
-    return false
-
-  for peerId in toQuery:
-    state.attempts[peerId] = state.attempts.getOrDefault(peerId, 0) + 1
-
-  debug "Lookup queries", peersToQuery = toQuery.mapIt(it.shortLog())
-
+type DispatchOutcome = enum
+  Completed
+  Errored
+
+type DispatchResult = object
+  peer: PeerId
+  outcome: DispatchOutcome
+  msg: Message
+
+type RpcFuture = Future[DispatchResult].Raising([CancelledError])
+
+type Attempt = object
+  peer: PeerId
+  fut: RpcFuture
+  deadline: Moment
+  abandoned: bool
+    ## its ``timeout`` elapsed; the slot is freed but the RPC
+    ## keeps running so it can still deliver, and its late result is ignored.
+
+proc runDispatch(
+    kad: KadDHT, peerId: PeerId, target: Key, dispatch: DispatchProc
+): Future[DispatchResult] {.async: (raises: [CancelledError]).} =
+  let res = await dispatch(kad, peerId, target)
+  if res.isErr():
+    error "Kad lookup: RPC error", peer = peerId.shortLog(), msg = res.error()
+    return DispatchResult(peer: peerId, outcome: Errored)
+  return DispatchResult(peer: peerId, outcome: Completed, msg: res.value())
+
+func activePeers(pending: seq[Attempt]): HashSet[PeerId] {.raises: [].} =
+  var peers = initHashSet[PeerId]()
+  for a in pending:
+    if not a.abandoned:
+      peers.incl(a.peer)
+  peers
+
+proc refill(
+    kad: KadDHT, state: LookupState, pending: var seq[Attempt], dispatch: DispatchProc
+) {.raises: [].} =
+  ## Keep up to ``alpha`` RPCs in flight by dispatching the next-closest
+  ## not-yet-active peers into any free slots.
+  var active = pending.activePeers()
   let target = state.target
-  let dispatchWithPeer = proc(
-      peerId: PeerId
-  ): Future[(PeerId, Result[Message, string])] {.
-      async: (raises: [CancelledError]), gcsafe
-  .} =
-    let msg = await dispatch(kad, peerId, target)
-    return (peerId, msg)
-
-  let rpcBatch = toQuery.mapIt(dispatchWithPeer(it))
-  for (fut, peerId) in zip(rpcBatch, toQuery):
-    state.inflight.mgetOrPut(peerId, @[]).add(FutureBase(fut))
-  let completedRPCBatch = await rpcBatch.collectCompleted(kad.config.timeout)
-
-  for (fut, peerId) in zip(rpcBatch, toQuery):
-    if not fut.finished() or fut.cancelled():
-      continue
-    if fut.failed():
-      state.responded[peerId] = RespondedStatus.Failed
-    else:
-      let (_, res) = fut.value()
-      if res.isErr():
-        state.responded[peerId] = RespondedStatus.Failed
-        error "Kad lookup: RPC error", peer = peerId.shortLog(), msg = res.error()
-      else:
-        state.responded[peerId] = RespondedStatus.Success
-
-  var toCancel: seq[FutureBase]
-  for (peerId, res) in completedRPCBatch:
-    let reply = res.valueOr:
+  for (peerId, _) in state.sortedShortlist():
+    if active.len >= kad.config.alpha:
+      break
+    if peerId in active:
       continue
-    let newPeerInfos = state.updateShortlist(reply)
-    kad.switch.updatePeers(
-      kad.config.addressPolicy, rtable, newPeerInfos, kad.config.limits.maxPeersPerIp,
-      kad.config.limits.maxPeersPerIpv4Subnet, kad.config.limits.maxPeersPerIpv6Subnet,
+    state.attempts[peerId] = state.attempts.getOrDefault(peerId, 0) + 1
+    debug "Lookup query", peer = peerId.shortLog()
+    pending.add(
+      Attempt(
+        peer: peerId,
+        fut: kad.runDispatch(peerId, target, dispatch),
+        deadline: Moment.now() + kad.config.timeout,
+        abandoned: false,
+      )
     )
-    await onReply(peerId, Opt.some(reply), state)
-
-  # Evicted peers are no longer eligible for retries, so cancel any abandoned RPCs.
-  for peerId in state.inflight.keys.toSeq:
-    if not state.shortlist.hasKey(peerId):
-      toCancel.add(state.inflight.getOrDefault(peerId).filterIt(not it.finished()))
-      state.inflight.del(peerId)
-
-  for peerId in toQuery:
-    if state.responded.hasKey(peerId) or
-        state.attempts.getOrDefault(peerId, 0) > kad.config.retries:
-      toCancel.add(state.inflight.getOrDefault(peerId).filterIt(not it.finished()))
-      state.inflight.del(peerId)
-  await toCancel.cancelAndWait()
+    active.incl(peerId)
+
+proc awaitProgress(pending: seq[Attempt]) {.async: (raises: [CancelledError]).} =
+  ## Wake as soon as any in-flight RPC finishes or the earliest active slot's
+  ## ``timeout`` elapses, whichever comes first.
+  var earliest = Opt.none(Moment)
+  for a in pending:
+    if not a.abandoned and (earliest.isNone or a.deadline < earliest.get()):
+      earliest = Opt.some(a.deadline)
+
+  let timer = sleepAsync(
+    if earliest.isSome:
+      max(earliest.get() - Moment.now(), ZeroDuration)
+    else:
+      InfiniteDuration
+  )
+  defer:
+    timer.cancelSoon()
 
-  return true
+  var futs = pending.mapIt(FutureBase(it.fut))
+  futs.add(FutureBase(timer))
+  try:
+    discard await race(futs)
+  except ValueError:
+    raiseAssert "race() cannot raise ValueError on a non-empty future list"
 
 proc iterativeLookup*(
     kad: KadDHT,
@@ -446,17 +456,71 @@ proc iterativeLookup*(
     onReply: ReplyHandler,
     stopCond: StopCond,
 ): Future[LookupState] {.async: (raises: [CancelledError]).} =
+  ## Drive an iterative lookup as a continuous ``alpha``-concurrency pipeline:
+  ## keep up to ``alpha`` RPCs in flight and re-dispatch the next-closest
+  ## unqueried peer as soon as any slot frees, so wall-clock tracks throughput
+  ## rather than the slowest peer in a synchronized round. A slow peer only
+  ## stalls its own slot for ``timeout`` before it is abandoned and retried.
   let state = LookupState.init(kad, target)
+  var pending: seq[Attempt]
 
-  while not stopCond(state):
-    if not await kad.lookOnce(state, rtable, dispatch, onReply):
-      break
+  defer:
+    await pending.mapIt(it.fut).cancelAndWait()
+
+  while true:
+    # Resolve finished RPCs; time out overdue ones by freeing their slot while
+    # leaving the RPC running (its late reply is ignored).
+    let now = Moment.now()
+    var completed: seq[DispatchResult]
+    var stillPending: seq[Attempt]
+    for a in pending:
+      if a.fut.finished():
+        if not a.abandoned and not a.fut.cancelled():
+          completed.add(a.fut.value())
+        continue
+      if not a.abandoned and now >= a.deadline:
+        stillPending.add(
+          Attempt(peer: a.peer, fut: a.fut, deadline: a.deadline, abandoned: true)
+        )
+      else:
+        stillPending.add(a)
+    pending = stillPending
+
+    var toCancel: seq[RpcFuture]
+    for res in completed:
+      case res.outcome
+      of Errored:
+        state.responded[res.peer] = RespondedStatus.Failed
+      of Completed:
+        state.responded[res.peer] = RespondedStatus.Success
+        let newPeerInfos = state.updateShortlist(res.msg)
+        kad.switch.updatePeers(
+          kad.config.addressPolicy, rtable, newPeerInfos,
+          kad.config.limits.maxPeersPerIp, kad.config.limits.maxPeersPerIpv4Subnet,
+          kad.config.limits.maxPeersPerIpv6Subnet,
+        )
+        await onReply(res.peer, Opt.some(res.msg), state)
+
+    # A responded peer needs no duplicate retries, and a newly discovered closer
+    # peer can evict an in-flight target from the shortlist; cancel both.
+    var keep: seq[Attempt]
+    for a in pending:
+      if state.responded.hasKey(a.peer) or not state.shortlist.hasKey(a.peer):
+        toCancel.add(a.fut)
+      else:
+        keep.add(a)
+    pending = keep
+    await toCancel.cancelAndWait()
+
+    # Once the stop condition holds, dispatch no new peers but keep draining the
+    # replies already in flight, mirroring the old round that finished its batch
+    # before terminating.
+    if not stopCond(state):
+      kad.refill(state, pending, dispatch)
 
-  var leftover: seq[FutureBase]
-  for futs in state.inflight.values:
-    leftover.add(futs.filterIt(not it.finished()))
-  if leftover.len > 0:
-    await leftover.cancelAndWait()
+    if pending.activePeers().len == 0:
+      break
+    await awaitProgress(pending)
 
   state
 
diff --git a/libp2p/protocols/service_discovery/random_find.nim b/libp2p/protocols/service_discovery/random_find.nim
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0 OR MIT
 # Copyright (c) Status Research & Development GmbH
 
-import std/[sequtils, sets, tables]
+import std/[sequtils, sets]
 import chronos, chronicles, results
 import ../../[peerid, peerinfo, switch, multihash, routing_record, extended_peer_record]
 import ../protocol
@@ -22,23 +22,22 @@ proc randomRecords(
 
   let randomKey = randomPeerId.toKey()
 
-  let state = LookupState.init(disco, randomKey)
   var queried: HashSet[PeerId]
   var getValFuts: seq[Future[Result[Message, string]].Raising([CancelledError])]
 
+  # Issue a getValue as soon as a peer responds, so they run in parallel with
+  # the rest of the lookup.
+  let onReply = proc(
+      peerId: PeerId, msg: Opt[Message], state: LookupState
+  ): Future[void] {.async: (raises: []), gcsafe.} =
+    if peerId notin queried:
+      queried.incl(peerId)
+      getValFuts.add(disco.dispatchGetVal(peerId, peerId.toKey()))
+
   try:
-    while not closestAvailableStop(state):
-      let progressed =
-        await disco.lookOnce(state, disco.rtable, findNodeDispatch, noopReply)
-      if not progressed:
-        break
-
-      # Issue getValue requests as soon as peers respond, so they run in
-      # parallel with the next lookup round.
-      for peerId, status in state.responded:
-        if status == RespondedStatus.Success and peerId notin queried:
-          queried.incl(peerId)
-          getValFuts.add(disco.dispatchGetVal(peerId, peerId.toKey()))
+    discard await disco.iterativeLookup(
+      randomKey, findNodeDispatch, onReply, closestAvailableStop
+    )
   except CancelledError as e:
     await noCancel allFutures(getValFuts.mapIt(it.cancelAndWait()))
     raise e
diff --git a/tests/libp2p/service_discovery/component/test_find_random.nim b/tests/libp2p/service_discovery/component/test_find_random.nim
@@ -26,7 +26,7 @@ suite "Service Discovery Component - Find Random":
       check record.peerId in peerIds
 
   asyncTest "lookupRandom completes when there are no peers to query":
-    # With an empty routing table the shortlist is empty, so lookOnce returns
+    # With an empty routing table the shortlist is empty, so the lookup returns
     # immediately and lookupRandom must complete rather than hang.
     let discos = setupServiceDiscoveryNodes(1)
     startAndDeferStop(discos)
