fix(p2p): better batch connection sampling (#13674)

## Overview

Improves batch connection sampling such that we can increase the retry
attempts for prover nodes requesting
transaction information from their peers.

### Core improvements
**Previously**
When making a batch request, we requested the peer list, iterated
through it until we found a peer without an reqresp active connection
and used those. This meant that we were commonly only using the first
3-4 peers in the peer list.

Now it will do random sampling without replacement to peers without
already active reqresp connections. If there are non available, then we
will just sample from peers already servicing requests.

Author: Maddiaa0

yarn-project/p2p/src/services/reqresp/connection-sampler/batch_connection_sampler.test.ts

@@ -22,7 +22,7 @@ describe('BatchConnectionSampler', () => {
 
     // Mock libp2p to return our test peers
     libp2p = {
-      getPeers: jest.fn().mockReturnValue(peers),
+      getPeers: jest.fn().mockImplementation(() => [...peers]),
     } as unknown as jest.Mocked<Libp2p>;
 
     // Create a real connection sampler with mocked random sampling
@@ -45,8 +45,7 @@ describe('BatchConnectionSampler', () => {
 
   it('assigns requests to peers deterministically with wraparound', () => {
     // Mock to return first two peers
-    let callCount = 0;
-    mockRandomSampler.random.mockImplementation(() => callCount++ % 2);
+    mockRandomSampler.random.mockImplementation(() => 0);
 
     // With 5 requests and 2 peers:
     // floor(5/2) = 2 requests per peer
@@ -66,9 +65,7 @@ describe('BatchConnectionSampler', () => {
   });
 
   it('handles peer removal and replacement', () => {
-    mockRandomSampler.random.mockImplementation(_ => {
-      return 2; // Return index 2 for replacement peer
-    });
+    mockRandomSampler.random.mockImplementation(_ => 0);
 
     // With 4 requests and 2 peers:
     // floor(4/2) = 2 requests per peer
@@ -80,6 +77,8 @@ describe('BatchConnectionSampler', () => {
     const initialPeer = sampler.getPeerForRequest(0);
     expect(initialPeer).toBe(peers[0]);
 
+    // Mock random to return the third peer
+    mockRandomSampler.random.mockImplementation(_ => 2);
     sampler.removePeerAndReplace(peers[0]);
 
     // After replacement:
@@ -94,7 +93,7 @@ describe('BatchConnectionSampler', () => {
   });
 
   it('handles peer removal and replacement - no replacement available', () => {
-    mockRandomSampler.random.mockImplementation(() => 2);
+    mockRandomSampler.random.mockImplementation(() => 0);
     const sampler = new BatchConnectionSampler(connectionSampler, /* batchSize */ 4, /* maxPeers */ 2);
 
     expect(sampler.activePeerCount).toBe(2);
@@ -112,13 +111,7 @@ describe('BatchConnectionSampler', () => {
   });
 
   it('distributes requests according to documentation example', () => {
-    let callCount = 0;
-    mockRandomSampler.random.mockImplementation(() => {
-      if (callCount < 3) {
-        return callCount++;
-      }
-      return 0;
-    });
+    mockRandomSampler.random.mockImplementation(() => 0);
 
     // Example from doc comment:
     // Peers:    [P1]      [P2]     [P3]
@@ -146,8 +139,7 @@ describe('BatchConnectionSampler', () => {
   });
 
   it('same number of requests per peers', () => {
-    let callCount = 0;
-    mockRandomSampler.random.mockImplementation(() => callCount++ % 2);
+    mockRandomSampler.random.mockImplementation(() => 0);
 
     const sampler = new BatchConnectionSampler(connectionSampler, /* batchSize */ 2, /* maxPeers */ 2);
     expect(sampler.requestsPerBucket).toBe(1);
@@ -165,10 +157,9 @@ describe('BatchConnectionSampler', () => {
     expect(sampler.activePeerCount).toBe(0);
     expect(sampler.getPeerForRequest(0)).toBeUndefined();
 
-    let i = 0;
-    mockRandomSampler.random.mockImplementation(() => i++ % 3);
+    mockRandomSampler.random.mockImplementation(() => 0);
 
-    libp2p.getPeers.mockReturnValue(peers);
+    libp2p.getPeers.mockImplementation(() => [...peers]);
     const samplerWithMorePeers = new BatchConnectionSampler(connectionSampler, /* batchSize */ 2, /* maxPeers */ 3);
     expect(samplerWithMorePeers.requestsPerBucket).toBe(1); // floor(2/3) = 0
     // First two requests go to first two peers

yarn-project/p2p/src/services/reqresp/connection-sampler/batch_connection_sampler.ts

@@ -70,11 +70,11 @@ export class BatchConnectionSampler {
 
     if (newPeer) {
       this.batch[index] = newPeer;
-      this.logger.trace(`Replaced peer ${peerId} with ${newPeer}`, { peerId, newPeer });
+      this.logger.trace('Replaced peer', { peerId, newPeer });
     } else {
       // If we couldn't get a replacement, remove the peer and compact the array
       this.batch.splice(index, 1);
-      this.logger.trace(`Removed peer ${peerId}`, { peerId });
+      this.logger.trace('Removed peer', { peerId });
     }
   }
 

yarn-project/p2p/src/services/reqresp/connection-sampler/connection_sampler.test.ts

@@ -20,7 +20,7 @@ describe('ConnectionSampler', () => {
 
     // Mock libp2p
     mockLibp2p = {
-      getPeers: jest.fn().mockReturnValue([...peers]),
+      getPeers: jest.fn().mockImplementation(() => [...peers]),
       dialProtocol: jest.fn(),
     };
 
@@ -196,15 +196,28 @@ describe('ConnectionSampler', () => {
 
       // Mock libp2p
       mockLibp2p = {
-        getPeers: jest.fn().mockReturnValue(peers),
+        getPeers: jest.fn().mockImplementation(() => [...peers]),
         dialProtocol: jest.fn(),
       };
 
       mockRandomSampler = mock<RandomSampler>();
+      mockRandomSampler.random.mockReturnValue(0);
       sampler = new ConnectionSampler(mockLibp2p, 1000, mockRandomSampler);
     });
 
+    it('should only return samples as many peers as available', () => {
+      const sampledPeers = sampler.samplePeersBatch(100);
+
+      expect(sampledPeers).toHaveLength(peers.length);
+    });
+
     it('prioritizes peers without active connections', () => {
+      mockRandomSampler.random
+        // Will pick the peers with active connections
+        .mockReturnValueOnce(3)
+        .mockReturnValueOnce(3)
+        .mockReturnValue(0);
+
       // Set up some peers with active connections
       sampler['activeConnectionsCount'].set(peers[3], 1);
       sampler['activeConnectionsCount'].set(peers[4], 2);
@@ -248,7 +261,8 @@ describe('ConnectionSampler', () => {
       const sampledPeers = sampler.samplePeersBatch(3);
 
       expect(sampledPeers).toHaveLength(3);
-      expect(sampledPeers).toEqual(expect.arrayContaining([peers[0], peers[1], peers[2]]));
+      // The last one will be picked first, then the first one, then the second one
+      expect(sampledPeers).toEqual(expect.arrayContaining([peers[4], peers[0], peers[1]]));
     });
 
     it('handles case when fewer peers available than requested', () => {

yarn-project/p2p/src/services/reqresp/connection-sampler/connection_sampler.ts

@@ -17,7 +17,7 @@ export class RandomSampler {
 }
 
 /**
- * A class that samples peers from the libp2p node and returns a peer that we don't already have a connection open to.
+ * A class that samples peers from the libp2p node and returns a peer that we don't already have a reqresp connection open to.
  * If we already have a connection open, we try to sample a different peer.
  * We do this MAX_SAMPLE_ATTEMPTS times, if we still don't find a peer we just go for it.
  *
@@ -69,32 +69,65 @@ export class ConnectionSampler {
   getPeer(excluding?: Map<string, boolean>): PeerId | undefined {
     // In libp2p getPeers performs a shallow copy, so this array can be sliced from safetly
     const peers = this.libp2p.getPeers();
+    const { peer } = this.getPeerFromList(peers, excluding);
+    return peer;
+  }
 
+  /**
+   * Samples a peer from a list of peers, excluding those that have active (reqresp) connections or are in the exclusion list
+   *
+   * @param peers - The list of peers to sample from
+   * @param excluding - The peers to exclude from the sampling
+   * @returns - A peer from the list, or undefined if no peers are available,
+   *          - a boolean indicating if the peer has active connections, and
+   *          - all sampled peers - to enable optional resampling
+   *
+   * @dev The provided list peers, should be mutated by this function. This allows batch sampling
+   *      to be performed without making extra copies of the list.
+   */
+  getPeerFromList(
+    peers: PeerId[],
+    excluding?: Map<string, boolean>,
+  ): {
+    peer: PeerId | undefined;
+    sampledPeers: PeerId[];
+  } {
     if (peers.length === 0) {
-      return undefined;
+      return { peer: undefined, sampledPeers: [] };
     }
 
-    let randomIndex = this.sampler.random(peers.length);
-    let attempts = 0;
-
-    // Keep sampling while:
-    // - we haven't exceeded max attempts AND
-    // - either the peer has active connections OR is in the exclusion list
-    while (
-      attempts < MAX_SAMPLE_ATTEMPTS &&
-      ((this.activeConnectionsCount.get(peers[randomIndex]) ?? 0) > 0 ||
-        (excluding?.get(peers[randomIndex]?.toString()) ?? false))
-    ) {
+    const sampledPeers: PeerId[] = [];
+    // Try to find a peer that has no active connections and is not in the exclusion list
+    for (let attempts = 0; attempts < MAX_SAMPLE_ATTEMPTS && peers.length > 0; attempts++) {
+      const randomIndex = this.sampler.random(peers.length);
+      const peer = peers[randomIndex];
+      const hasActiveConnections = (this.activeConnectionsCount.get(peer) ?? 0) > 0;
+      const isExcluded = excluding?.get(peer.toString()) ?? false;
+
+      // Remove this peer from consideration
       peers.splice(randomIndex, 1);
-      randomIndex = this.sampler.random(peers.length);
-      attempts++;
+
+      // If peer is suitable (no active connections and not excluded), return it
+      if (!hasActiveConnections && !isExcluded) {
+        this.logger.trace('Sampled peer', {
+          attempts,
+          peer,
+        });
+        return { peer, sampledPeers };
+      }
+
+      // Keep track of peers that have active reqresp channels, batch sampling will use these to resample
+      sampledPeers.push(peer);
     }
 
-    this.logger.trace(`Sampled peer in ${attempts} attempts`, {
-      attempts,
-      peer: peers[randomIndex]?.toString(),
+    // If we've exhausted our attempts or peers list is empty, return the last peer if available
+    const lastPeer = peers.length > 0 ? peers[this.sampler.random(peers.length)] : undefined;
+
+    this.logger.trace('Sampled peer', {
+      attempts: MAX_SAMPLE_ATTEMPTS,
+      peer: lastPeer?.toString(),
     });
-    return peers[randomIndex];
+    return { peer: lastPeer, sampledPeers };
   }
 
   /**
@@ -105,34 +138,41 @@ export class ConnectionSampler {
    */
   samplePeersBatch(numberToSample: number): PeerId[] {
     const peers = this.libp2p.getPeers();
-    const sampledPeers: PeerId[] = [];
-    const peersWithConnections: PeerId[] = []; // Hold onto peers with active connections incase we need to sample more
+    this.logger.debug('Sampling peers batch', { numberToSample, peers });
 
-    for (const peer of peers) {
-      const activeConnections = this.activeConnectionsCount.get(peer) ?? 0;
-      if (activeConnections === 0) {
-        if (sampledPeers.push(peer) === numberToSample) {
-          return sampledPeers;
-        }
-      } else {
-        peersWithConnections.push(peer);
+    // Only sample as many peers as we have available
+    numberToSample = Math.min(numberToSample, peers.length);
+
+    const batch: PeerId[] = [];
+    const withActiveConnections: Set<PeerId> = new Set();
+    for (let i = 0; i < numberToSample; i++) {
+      const { peer, sampledPeers } = this.getPeerFromList(peers, undefined);
+      if (peer) {
+        batch.push(peer);
+      }
+      if (sampledPeers.length > 0) {
+        sampledPeers.forEach(peer => withActiveConnections.add(peer));
       }
     }
+    const lengthWithoutConnections = batch.length;
 
     // If we still need more peers, sample from those with connections
-    while (sampledPeers.length < numberToSample && peersWithConnections.length > 0) {
-      const randomIndex = this.sampler.random(peersWithConnections.length);
-      const [peer] = peersWithConnections.splice(randomIndex, 1);
-      sampledPeers.push(peer);
+    while (batch.length < numberToSample && withActiveConnections.size > 0) {
+      const randomIndex = this.sampler.random(withActiveConnections.size);
+
+      const peer = Array.from(withActiveConnections)[randomIndex];
+      withActiveConnections.delete(peer);
+      batch.push(peer);
     }
 
-    this.logger.trace(`Batch sampled ${sampledPeers.length} unique peers`, {
-      peers: sampledPeers,
-      withoutConnections: sampledPeers.length - peersWithConnections.length,
-      withConnections: peersWithConnections.length,
+    this.logger.trace('Batch sampled peers', {
+      length: batch.length,
+      peers: batch,
+      withoutConnections: lengthWithoutConnections,
+      withConnections: numberToSample - lengthWithoutConnections,
     });
 
-    return sampledPeers;
+    return batch;
   }
 
   // Set of passthrough functions to keep track of active connections
@@ -155,8 +195,9 @@ export class ConnectionSampler {
     const updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerId) ?? 0) + 1;
     this.activeConnectionsCount.set(peerId, updatedActiveConnectionsCount);
 
-    this.logger.trace(`Dialed protocol ${protocol} with peer ${peerId.toString()}`, {
+    this.logger.trace('Dialed protocol', {
       streamId: stream.id,
+      protocol,
       peerId: peerId.toString(),
       activeConnectionsCount: updatedActiveConnectionsCount,
     });
@@ -181,7 +222,7 @@ export class ConnectionSampler {
       const updatedActiveConnectionsCount = (this.activeConnectionsCount.get(peerId) ?? 1) - 1;
       this.activeConnectionsCount.set(peerId, updatedActiveConnectionsCount);
 
-      this.logger.trace(`Closing connection to peer ${peerId.toString()}`, {
+      this.logger.trace('Closing connection', {
         streamId,
         peerId: peerId.toString(),
         protocol: stream.protocol,
@@ -207,7 +248,7 @@ export class ConnectionSampler {
         // Check if we have lost track of accounting
         if (this.activeConnectionsCount.get(peerId) === 0) {
           await this.close(streamId);
-          this.logger.debug(`Cleaned up stale connection ${streamId} to peer ${peerId.toString()}`);
+          this.logger.debug('Cleaned up stale connection', { streamId, peerId: peerId.toString() });
         }
       } catch (error) {
         this.logger.error(`Error cleaning up stale connection ${streamId}`, { error });

yarn-project/p2p/src/services/reqresp/reqresp.ts

@@ -196,7 +196,7 @@ export class ReqResp {
         this.logger.trace(`Sending request to peer: ${peer.toString()}`);
         const response = await this.sendRequestToPeer(peer, subProtocol, requestBuffer);
 
-        if (response && response.status !== ReqRespStatus.SUCCESS) {
+        if (response.status !== ReqRespStatus.SUCCESS) {
           this.logger.debug(
             `Request to peer ${peer.toString()} failed with status ${prettyPrintReqRespStatus(response.status)}`,
           );
@@ -325,7 +325,7 @@ export class ReqResp {
                 const response = await this.sendRequestToPeer(peer, subProtocol, requestBuffers[index]);
 
                 // Check the status of the response buffer
-                if (response && response.status !== ReqRespStatus.SUCCESS) {
+                if (response.status !== ReqRespStatus.SUCCESS) {
                   this.logger.debug(
                     `Request to peer ${peer.toString()} failed with status ${prettyPrintReqRespStatus(
                       response.status,
@@ -421,7 +421,7 @@ export class ReqResp {
     peerId: PeerId,
     subProtocol: ReqRespSubProtocol,
     payload: Buffer,
-  ): Promise<ReqRespResponse | undefined> {
+  ): Promise<ReqRespResponse> {
     let stream: Stream | undefined;
     try {
       this.metrics.recordRequestSent(subProtocol);
@@ -439,6 +439,12 @@ export class ReqResp {
     } catch (e: any) {
       this.metrics.recordRequestError(subProtocol);
       this.handleResponseError(e, peerId, subProtocol);
+
+      // If there is an exception, we return an unknown response
+      return {
+        status: ReqRespStatus.FAILURE,
+        data: Buffer.from([]),
+      };
     } finally {
       // Only close the stream if we created it
       if (stream) {

yarn-project/p2p/src/services/reqresp/status.ts

@@ -5,6 +5,7 @@ export enum ReqRespStatus {
   SUCCESS = 0,
   RATE_LIMIT_EXCEEDED = 1,
   BADLY_FORMED_REQUEST = 2,
+  FAILURE = 126,
   UNKNOWN = 127,
 }
 
@@ -53,6 +54,8 @@ export function prettyPrintReqRespStatus(status: ReqRespStatus) {
       return 'RATE_LIMIT_EXCEEDED';
     case ReqRespStatus.BADLY_FORMED_REQUEST:
       return 'BADLY_FORMED_REQUEST';
+    case ReqRespStatus.FAILURE:
+      return 'FAILURE';
     case ReqRespStatus.UNKNOWN:
       return 'UNKNOWN';
   }