fix(kad-dht): prefer connected peers to reduce unnecessary dials

Author: paschal533

packages/kad-dht/src/content-routing/index.ts

@@ -151,16 +151,24 @@ export class ContentRouting {
         finalPeerEvents.push(event)
       }
 
-      finalPeerEvents.forEach(event => {
-        queue.add(async () => {
-          for await (const notifyEvent of publishProviderRecord(event)) {
-            events.push(notifyEvent)
-          }
+      // sort connected peers first to reuse existing connections before
+      // opening new ones when sending ADD_PROVIDER records
+      finalPeerEvents
+        .sort((a, b) => {
+          const aConnected = (self.components.connectionManager.getConnections(a.peer.id) ?? []).length > 0 ? 0 : 1
+          const bConnected = (self.components.connectionManager.getConnections(b.peer.id) ?? []).length > 0 ? 0 : 1
+          return aConnected - bConnected
         })
-          .catch(err => {
-            this.log.error('error publishing provider record to peer - %e', err)
+        .forEach(event => {
+          queue.add(async () => {
+            for await (const notifyEvent of publishProviderRecord(event)) {
+              events.push(notifyEvent)
+            }
           })
-      })
+            .catch(err => {
+              this.log.error('error publishing provider record to peer - %e', err)
+            })
+        })
     })
       .catch(err => {
         events.end(err)

packages/kad-dht/src/query/manager.ts

@@ -172,13 +172,19 @@ export class QueryManager implements Startable {
         count: this.routingTable.kBucketSize
       })
 
-      // split peers into d buckets evenly(ish)
-      const peersToQuery = peers.sort(() => {
-        if (Math.random() > 0.5) {
-          return 1
+      // split peers into d buckets evenly(ish), preferring already-connected
+      // peers to avoid opening new connections unnecessarily
+      const peersToQuery = peers.sort((a, b) => {
+        const aConnected = (this.connectionManager.getConnections(a) ?? []).length > 0 ? 0 : 1
+        const bConnected = (this.connectionManager.getConnections(b) ?? []).length > 0 ? 0 : 1
+
+        if (aConnected !== bConnected) {
+          // put connected peers first
+          return aConnected - bConnected
         }
 
-        return -1
+        // randomize within each group to maintain path diversity
+        return Math.random() > 0.5 ? 1 : -1
       })
         .reduce((acc: PeerId[][], curr, index) => {
           acc[index % this.disjointPaths].push(curr)

packages/kad-dht/test/content-routing.spec.ts

@@ -354,4 +354,34 @@ describe('content routing', () => {
     expect(sendMessageSpy.callCount).to.equal(initialMessageCalls,
       'No new network calls should be made after abort')
   })
+
+  it('queries connection status to prioritise connected peers when sending ADD_PROVIDER', async function () {
+    this.timeout(20 * 1000)
+
+    const dhts = await sortDHTs(await Promise.all([
+      testDHT.spawn(),
+      testDHT.spawn(),
+      testDHT.spawn(),
+      testDHT.spawn()
+    ]), await kadUtils.convertBuffer(cid.multihash.bytes))
+
+    // connect all peers to the provider (dhts[3]) so they appear in its routing table
+    await Promise.all([
+      testDHT.connect(dhts[0], dhts[3]),
+      testDHT.connect(dhts[1], dhts[3]),
+      testDHT.connect(dhts[2], dhts[3])
+    ])
+
+    // getConnections is already a sinon stub from stubInterface — reset its history
+    // and verify it gets called during provide() for the connected-peers-first sort
+    dhts[3].components.connectionManager.getConnections.resetHistory()
+
+    await drain(dhts[3].dht.provide(cid))
+
+    // getConnections should have been called to check peer connection status
+    // when sorting the closest peers for ADD_PROVIDER fan-out
+    expect(dhts[3].components.connectionManager.getConnections.called).to.be.true(
+      'getConnections should be called to prioritise connected peers during provide()'
+    )
+  })
 })

packages/kad-dht/test/query.spec.ts

@@ -998,4 +998,63 @@ describe('QueryManager', () => {
 
     await manager.stop()
   })
+
+  it('should prefer connected peers when distributing to disjoint paths', async () => {
+    // pick peers at higher XOR indices as "connected" and lower-XOR as "disconnected"
+    // so XOR-distance sorting alone would put disconnected peers first — only our
+    // connected-first bucket assignment can ensure each path gets a connected peer
+    const connectedPeers = peers.slice(10, 12).map(p => p.peerId)
+    const disconnectedPeers = peers.slice(0, 2).map(p => p.peerId)
+    // routing table returns the 4 peers in arbitrary order
+    const startingPeers = [...connectedPeers, ...disconnectedPeers]
+
+    const connectionManager = stubInterface<ConnectionManager>({
+      isDialable: async () => true
+    })
+    connectionManager.getConnections.callsFake((peerId?: PeerId) => {
+      if (peerId != null && connectedPeers.some(p => p.equals(peerId))) {
+        return [{}] as any
+      }
+      return []
+    })
+
+    const manager = new QueryManager({
+      peerId: ourPeerId,
+      logger: defaultLogger(),
+      connectionManager
+    }, {
+      ...defaultInit(),
+      disjointPaths: 2,
+      alpha: 1
+    })
+    await manager.start()
+
+    routingTable.closestPeers.returns(startingPeers)
+
+    // track which peers are queried on each disjoint path
+    const peersByPath = new Map<number, PeerId[]>([[0, []], [1, []]])
+
+    const queryFunc: QueryFunc = async function * ({ peer, path }) {
+      peersByPath.get(path.index)?.push(peer.id)
+
+      yield peerResponseEvent({
+        from: peer.id,
+        messageType: MessageType.GET_VALUE,
+        path
+      })
+    }
+
+    await drain(manager.run(key, queryFunc))
+
+    // with connected-peers-first sort, even though disconnected peers have lower
+    // XOR distance, each disjoint path should get at least one connected starting peer
+    for (const [pathIndex, queriedPeers] of peersByPath) {
+      const hasConnectedPeer = queriedPeers.some(p => connectedPeers.some(c => c.equals(p)))
+      expect(hasConnectedPeer).to.be.true(
+        `path ${pathIndex} should have at least one connected starting peer`
+      )
+    }
+
+    await manager.stop()
+  })
 })