feat(reqresp): send status messages along with reqresp responses

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

@@ -3,6 +3,8 @@ import { Fr } from '@aztec/foundation/fields';
 
 import { type PeerId } from '@libp2p/interface';
 
+import { ReqRespStatus } from './status.js';
+
 /*
  * Request Response Sub Protocols
  */
@@ -31,6 +33,15 @@ export type ReqRespSubProtocolHandler = (peerId: PeerId, msg: Buffer) => Promise
  */
 export type ReqRespSubProtocolRateLimits = Record<ReqRespSubProtocol, ProtocolRateLimitQuota>;
 
+/**
+ * The response from the ReqResp protocol
+ * Consists of a status (Error code) and data
+ */
+export interface ReqRespResponse {
+  status: ReqRespStatus;
+  data: Buffer;
+}
+
 /**
  * A rate limit quota
  */

yarn-project/p2p/src/services/reqresp/rate-limiter/rate_limiter.ts

@@ -200,9 +200,11 @@ export class RequestResponseRateLimiter {
 
     switch (rateLimitStatus) {
       case RateLimitStatus.DeniedPeer:
-        this.peerScoring.penalizePeer(peerId, PeerErrorSeverity.MidToleranceError);
+        // Hitting a peer specific limit, we should lightly penalise the peer
+        this.peerScoring.penalizePeer(peerId, PeerErrorSeverity.LowToleranceError);
         return false;
       case RateLimitStatus.DeniedGlobal:
+        // Hitting a global limit, we should not penalise the peer
         return false;
       default:
         return true;

yarn-project/p2p/src/services/reqresp/rate-limiter/rate_limits.ts

@@ -25,11 +25,11 @@ export const DEFAULT_RATE_LIMITS: ReqRespSubProtocolRateLimits = {
   [ReqRespSubProtocol.TX]: {
     peerLimit: {
       quotaTimeMs: 1000,
-      quotaCount: 5,
+      quotaCount: 10,
     },
     globalLimit: {
       quotaTimeMs: 1000,
-      quotaCount: 10,
+      quotaCount: 20,
     },
   },
   [ReqRespSubProtocol.BLOCK]: {

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

@@ -22,6 +22,7 @@ import { type PeerScoring } from '../peer-manager/peer_scoring.js';
 import { ReqRespSubProtocol, RequestableBuffer } from './interface.js';
 import { reqRespBlockHandler } from './protocols/block.js';
 import { GoodByeReason, reqGoodbyeHandler } from './protocols/goodbye.js';
+import { ReqRespStatus, prettyPrintReqRespStatus } from './status.js';
 
 const PING_REQUEST = RequestableBuffer.fromBuffer(Buffer.from('ping'));
 
@@ -58,6 +59,7 @@ describe('ReqResp', () => {
     await sleep(500);
 
     const res = await pinger.sendRequest(ReqRespSubProtocol.PING, PING_REQUEST);
+    logger.info(`Res: ${res?.toBuffer().toString('utf-8')}`);
 
     await sleep(500);
     expect(res?.toBuffer().toString('utf-8')).toEqual('pong');
@@ -126,10 +128,21 @@ describe('ReqResp', () => {
     await sleep(500);
 
     // Default rate is set at 1 every 200 ms; so this should fire a few times
+    const responses = [];
     for (let i = 0; i < 10; i++) {
-      await nodes[0].req.sendRequestToPeer(nodes[1].p2p.peerId, ReqRespSubProtocol.PING, Buffer.from('ping'));
+      // Response object contains the status (error flags) and data
+      const response = await nodes[0].req.sendRequestToPeer(
+        nodes[1].p2p.peerId,
+        ReqRespSubProtocol.PING,
+        Buffer.from('ping'),
+      );
+      responses.push(response);
     }
 
+    // Check that one of the responses gets a rate limit response
+    const rateLimitResponse = responses.find(response => response?.status === ReqRespStatus.RATE_LIMIT_EXCEEDED);
+    expect(rateLimitResponse).toBeDefined();
+
     // Make sure the error message is logged
     const errorMessage = `Rate limit exceeded for ${ReqRespSubProtocol.PING} from ${nodes[0].p2p.peerId.toString()}`;
     expect(loggerSpy).toHaveBeenCalledWith(expect.stringContaining(errorMessage));
@@ -343,7 +356,8 @@ describe('ReqResp', () => {
       );
 
       // Expect the response to be a buffer of length 1
-      expect(response).toEqual(Buffer.from([0x0]));
+      expect(response?.status).toEqual(ReqRespStatus.SUCCESS);
+      expect(response?.data).toEqual(Buffer.from([0x0]));
     });
   });
 
@@ -413,6 +427,8 @@ describe('ReqResp', () => {
       const batchSize = 12;
       nodes = await createNodes(peerScoring, 3);
 
+      const requesterLoggerSpy = jest.spyOn((nodes[0].req as any).logger, 'debug');
+
       await startNodes(nodes);
       await sleep(500);
       await connectToPeers(nodes);
@@ -426,6 +442,11 @@ describe('ReqResp', () => {
 
       const res = await nodes[0].req.sendBatchRequest(ReqRespSubProtocol.PING, requests);
       expect(res).toEqual(expectResponses);
+
+      // Check that we did detect hitting a rate limit
+      expect(requesterLoggerSpy).toHaveBeenCalledWith(
+        expect.stringContaining(`${prettyPrintReqRespStatus(ReqRespStatus.RATE_LIMIT_EXCEEDED)}`),
+      );
     });
   });
 });

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

@@ -22,14 +22,16 @@ import { ConnectionSampler } from './connection-sampler/connection_sampler.js';
 import {
   DEFAULT_SUB_PROTOCOL_HANDLERS,
   DEFAULT_SUB_PROTOCOL_VALIDATORS,
-  type ReqRespSubProtocol,
+  ReqRespResponse,
+  ReqRespSubProtocol,
   type ReqRespSubProtocolHandlers,
   type ReqRespSubProtocolValidators,
   type SubProtocolMap,
   subProtocolMap,
 } from './interface.js';
 import { ReqRespMetrics } from './metrics.js';
 import { RequestResponseRateLimiter } from './rate-limiter/rate_limiter.js';
+import { ReqRespStatus, prettyPrintReqRespStatus } from './status.js';
 
 /**
  * The Request Response Service
@@ -190,13 +192,21 @@ 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) {
+          this.logger.debug(
+            `Request to peer ${peer.toString()} failed with status ${prettyPrintReqRespStatus(response.status)}`,
+          );
+          continue;
+        }
+
         // If we get a response, return it, otherwise we iterate onto the next peer
         // We do not consider it a success if we have an empty buffer
-        if (response && response.length > 0) {
-          const object = subProtocolMap[subProtocol].response.fromBuffer(response);
+        if (response && response.data.length > 0) {
+          const object = subProtocolMap[subProtocol].response.fromBuffer(response.data);
           // The response validator handles peer punishment within
           const isValid = await responseValidator(request, object, peer);
           if (!isValid) {
+            // TODO(md): return the invalid response message!
             throw new InvalidResponseError();
           }
           return object;
@@ -311,8 +321,22 @@ export class ReqResp {
               for (const index of indices) {
                 const response = await this.sendRequestToPeer(peer, subProtocol, requestBuffers[index]);
 
-                if (response && response.length > 0) {
-                  const object = subProtocolMap[subProtocol].response.fromBuffer(response);
+                // Check the status of the response buffer
+                if (response && response.status !== ReqRespStatus.SUCCESS) {
+                  this.logger.debug(
+                    `Request to peer ${peer.toString()} failed with status ${prettyPrintReqRespStatus(
+                      response.status,
+                    )}`,
+                  );
+
+                  // If we hit a rate limit or some failure, we remove the peer and return the results,
+                  // they will be split among remaining peers and the new sampled peer
+                  batchSampler.removePeerAndReplace(peer);
+                  return { peer, results: peerResults };
+                }
+
+                if (response && response.data.length > 0) {
+                  const object = subProtocolMap[subProtocol].response.fromBuffer(response.data);
                   const isValid = await responseValidator(requests[index], object, peer);
 
                   if (isValid) {
@@ -394,16 +418,16 @@ export class ReqResp {
     peerId: PeerId,
     subProtocol: ReqRespSubProtocol,
     payload: Buffer,
-  ): Promise<Buffer | undefined> {
+  ): Promise<ReqRespResponse | undefined> {
     let stream: Stream | undefined;
     try {
       this.metrics.recordRequestSent(subProtocol);
 
       stream = await this.connectionSampler.dialProtocol(peerId, subProtocol);
 
       // Open the stream with a timeout
-      const result = await executeTimeout<Buffer>(
-        (): Promise<Buffer> => pipe([payload], stream!, this.readMessage.bind(this)),
+      const result = await executeTimeout<ReqRespResponse>(
+        (): Promise<ReqRespResponse> => pipe([payload], stream!, this.readMessage.bind(this)),
         this.individualRequestTimeoutMs,
         () => new IndividualReqRespTimeoutError(),
       );
@@ -448,6 +472,14 @@ export class ReqResp {
    */
   private categorizeError(e: any, peerId: PeerId, subProtocol: ReqRespSubProtocol): PeerErrorSeverity | undefined {
     // Non pubishable errors
+    if (subProtocol === ReqRespSubProtocol.GOODBYE) {
+      this.logger.debug('Error encountered on goodbye sub protocol, no penalty', {
+        peerId: peerId.toString(),
+        subProtocol,
+      });
+      return undefined;
+    }
+
     // We do not punish a collective timeout, as the node triggers this interupt, independent of the peer's behaviour
     const logTags = {
       peerId: peerId.toString(),
@@ -493,13 +525,35 @@ export class ReqResp {
   /**
    * Read a message returned from a stream into a single buffer
    */
-  private async readMessage(source: AsyncIterable<Uint8ArrayList>): Promise<Buffer> {
+  private async readMessage(source: AsyncIterable<Uint8ArrayList>): Promise<ReqRespResponse> {
+    let statusBuffer: ReqRespStatus | undefined;
+
     const chunks: Uint8Array[] = [];
-    for await (const chunk of source) {
-      chunks.push(chunk.subarray());
+    try {
+      for await (const chunk of source) {
+        if (statusBuffer === undefined) {
+          const firstChunkBuffer = chunk.subarray();
+          statusBuffer = firstChunkBuffer[0];
+          this.logger.info(`First chunk: ${statusBuffer}`);
+        } else {
+          chunks.push(chunk.subarray());
+        }
+      }
+
+      const messageData = Buffer.concat(chunks);
+      const message: Buffer = this.snappyTransform.inboundTransformNoTopic(messageData);
+
+      return {
+        status: statusBuffer ?? ReqRespStatus.UNKNOWN,
+        data: message,
+      };
+    } catch (e: any) {
+      this.logger.error(`Reading message failed: ${e.message}`);
+      return {
+        status: ReqRespStatus.UNKNOWN,
+        data: Buffer.from([]),
+      };
     }
-    const messageData = Buffer.concat(chunks);
-    return this.snappyTransform.inboundTransformNoTopic(messageData);
   }
 
   /**
@@ -525,25 +579,38 @@ export class ReqResp {
   private async streamHandler(protocol: ReqRespSubProtocol, { stream, connection }: IncomingStreamData) {
     this.metrics.recordRequestReceived(protocol);
 
-    // Store a reference to from this for the async generator
-    if (!this.rateLimiter.allow(protocol, connection.remotePeer)) {
-      this.logger.warn(`Rate limit exceeded for ${protocol} from ${connection.remotePeer}`);
+    try {
+      // Store a reference to from this for the async generator
+      if (!this.rateLimiter.allow(protocol, connection.remotePeer)) {
+        this.logger.warn(`Rate limit exceeded for ${protocol} from ${connection.remotePeer}`);
 
-      // TODO(#8483): handle changing peer scoring for failed rate limit, maybe differentiate between global and peer limits here when punishing
-      await stream.close();
-      return;
-    }
+        // TODO(#8483): handle changing peer scoring for failed rate limit, maybe differentiate between global and peer limits here when punishing
 
-    const handler = this.subProtocolHandlers[protocol];
-    const transform = this.snappyTransform;
+        // TODO(md): make a generic error response chunk
+        throw new Error('Rate limit exceeded');
+
+        // await stream.close();
+        // return;
+      }
+
+      const handler = this.subProtocolHandlers[protocol];
+      const transform = this.snappyTransform;
+      const logger = this.logger;
 
-    try {
       await pipe(
         stream,
         async function* (source: any) {
           for await (const chunkList of source) {
             const msg = Buffer.from(chunkList.subarray());
             const response = await handler(connection.remotePeer, msg);
+
+            // Send success code first, then the response
+            const successChunk = Buffer.from([ReqRespStatus.SUCCESS]);
+            logger.info(`Success chunk: ${successChunk.toString('hex')}`);
+            yield new Uint8Array(successChunk);
+
+            logger.info(`Response: ${response.toString('hex')}`);
+
             yield new Uint8Array(transform.outboundTransformNoTopic(response));
           }
         },
@@ -552,8 +619,24 @@ export class ReqResp {
     } catch (e: any) {
       this.logger.warn(e);
       this.metrics.recordResponseError(protocol);
+
+      const sendErrorChunk = this.sendErrorChunk(ReqRespStatus.RATE_LIMIT_EXCEEDED);
+
+      // Return and yield the response chunk
+      await pipe(
+        stream,
+        async function* (_source: any) {
+          yield* sendErrorChunk;
+        },
+        stream,
+      );
     } finally {
       await stream.close();
     }
   }
+
+  private async *sendErrorChunk(error: ReqRespStatus): AsyncIterable<Uint8Array> {
+    const errorChunk = Buffer.from([error]);
+    yield new Uint8Array(errorChunk);
+  }
 }

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

@@ -0,0 +1,35 @@
+/**
+ * The error codes for the ReqResp protocol
+ */
+export enum ReqRespStatus {
+  SUCCESS = 0,
+  RATE_LIMIT_EXCEEDED = 1,
+
+  // TODO: update the below errors
+  COLLECTIVE_TIMEOUT = 2,
+  INDIVIDUAL_TIMEOUT = 3,
+  INVALID_RESPONSE = 4,
+  UNKNOWN = 5,
+}
+
+/**
+ * Pretty print the ReqResp status
+ * @param status
+ * @returns
+ */
+export function prettyPrintReqRespStatus(status: ReqRespStatus) {
+  switch (status) {
+    case ReqRespStatus.SUCCESS:
+      return 'SUCCESS';
+    case ReqRespStatus.RATE_LIMIT_EXCEEDED:
+      return 'RATE_LIMIT_EXCEEDED';
+    case ReqRespStatus.COLLECTIVE_TIMEOUT:
+      return 'COLLECTIVE_TIMEOUT';
+    case ReqRespStatus.INDIVIDUAL_TIMEOUT:
+      return 'INDIVIDUAL_TIMEOUT';
+    case ReqRespStatus.INVALID_RESPONSE:
+      return 'INVALID_RESPONSE';
+    case ReqRespStatus.UNKNOWN:
+      return 'UNKNOWN';
+  }
+}