feat: add insturmentation to attestation and epoch quote mem pools

yarn-project/circuit-types/src/p2p/block_attestation.ts

@@ -72,4 +72,8 @@ export class BlockAttestation extends Gossipable {
   static empty(): BlockAttestation {
     return new BlockAttestation(ConsensusPayload.empty(), Signature.empty());
   }
+
+  getSize(): number {
+    return this.payload.getSize() + this.signature.getSize();
+  }
 }

yarn-project/circuit-types/src/p2p/block_proposal.ts

@@ -75,4 +75,8 @@ export class BlockProposal extends Gossipable {
     const reader = BufferReader.asReader(buf);
     return new BlockProposal(reader.readObject(ConsensusPayload), reader.readObject(Signature));
   }
+
+  getSize(): number {
+    return this.payload.getSize() + this.signature.getSize();
+  }
 }

yarn-project/circuit-types/src/p2p/consensus_payload.ts

@@ -9,6 +9,10 @@ import { TxHash } from '../tx/tx_hash.js';
 import { type Signable } from './signature_utils.js';
 
 export class ConsensusPayload implements Signable {
+  // Cached values
+  private asBuffer: Buffer | undefined;
+  private size: number | undefined;
+
   constructor(
     /** The block header the attestation is made over */
     public readonly header: Header,
@@ -31,7 +35,12 @@ export class ConsensusPayload implements Signable {
   }
 
   toBuffer(): Buffer {
-    return serializeToBuffer([this.header, this.archive, this.txHashes.length, this.txHashes]);
+    if (this.asBuffer) {
+      return this.asBuffer;
+    }
+    this.asBuffer = serializeToBuffer([this.header, this.archive, this.txHashes.length, this.txHashes]);
+    this.size = this.asBuffer.length;
+    return this.asBuffer;
   }
 
   static fromBuffer(buf: Buffer | BufferReader): ConsensusPayload {
@@ -50,4 +59,17 @@ export class ConsensusPayload implements Signable {
   static empty(): ConsensusPayload {
     return new ConsensusPayload(Header.empty(), Fr.ZERO, []);
   }
+
+  /**
+   * Get the size of the consensus payload in bytes.
+   * @returns The size of the consensus payload.
+   */
+  getSize(): number {
+    // We cache size to avoid recalculating it
+    if (this.size) {
+      return this.size;
+    }
+    this.toBuffer();
+    return this.size!;
+  }
 }

yarn-project/circuit-types/src/p2p/gossipable.ts

@@ -23,4 +23,11 @@ export abstract class Gossipable {
    * - Serialization method
    */
   abstract toBuffer(): Buffer;
+
+  /**
+   * Get the size of the gossipable object.
+   *
+   * This is used for metrics recording.
+   */
+  abstract getSize(): number;
 }

yarn-project/circuit-types/src/prover_coordination/epoch_proof_quote.ts

@@ -69,4 +69,12 @@ export class EpochProofQuote extends Gossipable {
       signature: this.signature.toViemSignature(),
     };
   }
+
+  /**
+   * Get the size of the epoch proof quote in bytes.
+   * @returns The size of the epoch proof quote in bytes.
+   */
+  getSize(): number {
+    return this.payload.getSize() + this.signature.getSize();
+  }
 }

yarn-project/circuit-types/src/prover_coordination/epoch_proof_quote_payload.ts

@@ -5,6 +5,10 @@ import { type FieldsOf } from '@aztec/foundation/types';
 import { inspect } from 'util';
 
 export class EpochProofQuotePayload {
+  // Cached values
+  private asBuffer: Buffer | undefined;
+  private size: number | undefined;
+
   constructor(
     public readonly epochToProve: bigint,
     public readonly validUntilSlot: bigint,
@@ -24,7 +28,13 @@ export class EpochProofQuotePayload {
   }
 
   toBuffer(): Buffer {
-    return serializeToBuffer(...EpochProofQuotePayload.getFields(this));
+    // We cache the buffer to avoid recalculating it
+    if (this.asBuffer) {
+      return this.asBuffer;
+    }
+    this.asBuffer = serializeToBuffer(...EpochProofQuotePayload.getFields(this));
+    this.size = this.asBuffer.length;
+    return this.asBuffer;
   }
 
   static fromBuffer(buf: Buffer | BufferReader): EpochProofQuotePayload {
@@ -84,6 +94,16 @@ export class EpochProofQuotePayload {
     };
   }
 
+  getSize(): number {
+    // We cache size to avoid recalculating it
+    if (this.size) {
+      return this.size;
+    }
+    // Size is cached when calling toBuffer
+    this.toBuffer();
+    return this.size!;
+  }
+
   [inspect.custom](): string {
     return `EpochProofQuotePayload { epochToProve: ${this.epochToProve}, validUntilSlot: ${this.validUntilSlot}, bondAmount: ${this.bondAmount}, prover: ${this.prover}, basisPointFee: ${this.basisPointFee} }`;
   }

yarn-project/foundation/src/eth-signature/eth_signature.ts

@@ -19,6 +19,10 @@ export type ViemSignature = {
  * Contains a signature split into it's primary components (r,s,v)
  */
 export class Signature {
+  // Cached values
+  private asBuffer: Buffer | undefined;
+  private size: number | undefined;
+
   constructor(
     /** The r value of the signature */
     public readonly r: Buffer32,
@@ -73,7 +77,24 @@ export class Signature {
   }
 
   toBuffer(): Buffer {
-    return serializeToBuffer([this.r, this.s, this.v]);
+    // We cache the buffer to avoid recalculating it
+    if (this.asBuffer) {
+      return this.asBuffer;
+    }
+    this.asBuffer = serializeToBuffer([this.r, this.s, this.v]);
+    this.size = this.asBuffer.length;
+    return this.asBuffer;
+  }
+
+  getSize(): number {
+    // We cache size to avoid recalculating it
+    if (this.size) {
+      return this.size;
+    }
+
+    // Size is cached when calling toBuffer
+    this.toBuffer();
+    return this.size!;
   }
 
   to0xString(): `0x${string}` {

yarn-project/p2p/src/client/index.ts

@@ -38,8 +38,8 @@ export const createP2PClient = async (
 
   const mempools: MemPools = {
     txPool: deps.txPool ?? new AztecKVTxPool(store, telemetry),
-    attestationPool: deps.attestationPool ?? new InMemoryAttestationPool(),
-    epochProofQuotePool: deps.epochProofQuotePool ?? new MemoryEpochProofQuotePool(),
+    attestationPool: deps.attestationPool ?? new InMemoryAttestationPool(telemetry),
+    epochProofQuotePool: deps.epochProofQuotePool ?? new MemoryEpochProofQuotePool(telemetry),
   };
 
   let p2pService;

yarn-project/p2p/src/mem_pools/attestation_pool/memory_attestation_pool.test.ts

@@ -1,7 +1,11 @@
+import { type BlockAttestation } from '@aztec/circuit-types';
 import { Fr } from '@aztec/foundation/fields';
+import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
+import { type MockProxy, mock } from 'jest-mock-extended';
 import { type PrivateKeyAccount } from 'viem';
 
+import { type PoolInstrumentation } from '../instrumentation.js';
 import { InMemoryAttestationPool } from './memory_attestation_pool.js';
 import { generateAccount, mockAttestation } from './mocks.js';
 
@@ -10,10 +14,20 @@ const NUMBER_OF_SIGNERS_PER_TEST = 4;
 describe('MemoryAttestationPool', () => {
   let ap: InMemoryAttestationPool;
   let signers: PrivateKeyAccount[];
+  const telemetry = new NoopTelemetryClient();
+
+  // Check that metrics are recorded correctly
+  let metricsMock: MockProxy<PoolInstrumentation<BlockAttestation>>;
 
   beforeEach(() => {
-    ap = new InMemoryAttestationPool();
+    // Use noop telemetry client while testing.
+
+    ap = new InMemoryAttestationPool(telemetry);
     signers = Array.from({ length: NUMBER_OF_SIGNERS_PER_TEST }, generateAccount);
+
+    metricsMock = mock<PoolInstrumentation<BlockAttestation>>();
+    // Can i overwrite this like this??
+    (ap as any).metrics = metricsMock;
   });
 
   it('should add attestations to pool', async () => {
@@ -25,6 +39,9 @@ describe('MemoryAttestationPool', () => {
 
     await ap.addAttestations(attestations);
 
+    // Check metrics have been updated.
+    expect(metricsMock.recordAddedObjects).toHaveBeenCalledWith(attestations.length);
+
     const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
 
     expect(retreivedAttestations.length).toBe(NUMBER_OF_SIGNERS_PER_TEST);
@@ -33,6 +50,8 @@ describe('MemoryAttestationPool', () => {
     // Delete by slot
     await ap.deleteAttestationsForSlot(BigInt(slotNumber));
 
+    expect(metricsMock.recordRemovedObjects).toHaveBeenCalledWith(attestations.length);
+
     const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
     expect(retreivedAttestationsAfterDelete.length).toBe(0);
   });
@@ -82,12 +101,16 @@ describe('MemoryAttestationPool', () => {
 
     await ap.addAttestations(attestations);
 
+    expect(metricsMock.recordAddedObjects).toHaveBeenCalledWith(attestations.length);
+
     const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
     expect(retreivedAttestations.length).toBe(NUMBER_OF_SIGNERS_PER_TEST);
     expect(retreivedAttestations).toEqual(attestations);
 
     await ap.deleteAttestations(attestations);
 
+    expect(metricsMock.recordRemovedObjects).toHaveBeenCalledWith(attestations.length);
+
     const gottenAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
     expect(gottenAfterDelete.length).toBe(0);
   });
@@ -118,12 +141,16 @@ describe('MemoryAttestationPool', () => {
 
     await ap.addAttestations(attestations);
 
+    expect(metricsMock.recordAddedObjects).toHaveBeenCalledWith(attestations.length);
+
     const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
     expect(retreivedAttestations.length).toBe(NUMBER_OF_SIGNERS_PER_TEST);
     expect(retreivedAttestations).toEqual(attestations);
 
     await ap.deleteAttestationsForSlotAndProposal(BigInt(slotNumber), proposalId);
 
+    expect(metricsMock.recordRemovedObjects).toHaveBeenCalledWith(attestations.length);
+
     const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId);
     expect(retreivedAttestationsAfterDelete.length).toBe(0);
   });

yarn-project/p2p/src/mem_pools/attestation_pool/memory_attestation_pool.ts

@@ -1,13 +1,18 @@
 import { type BlockAttestation } from '@aztec/circuit-types';
 import { createDebugLogger } from '@aztec/foundation/log';
+import { type TelemetryClient } from '@aztec/telemetry-client';
 
+import { PoolInstrumentation } from '../instrumentation.js';
 import { type AttestationPool } from './attestation_pool.js';
 
 export class InMemoryAttestationPool implements AttestationPool {
+  private metrics: PoolInstrumentation<BlockAttestation>;
+
   private attestations: Map</*slot=*/ bigint, Map</*proposalId*/ string, Map</*address=*/ string, BlockAttestation>>>;
 
-  constructor(private log = createDebugLogger('aztec:attestation_pool')) {
+  constructor(telemetry: TelemetryClient, private log = createDebugLogger('aztec:attestation_pool')) {
     this.attestations = new Map();
+    this.metrics = new PoolInstrumentation(telemetry, 'InMemoryAttestationPool');
   }
 
   public getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
@@ -35,21 +40,46 @@ export class InMemoryAttestationPool implements AttestationPool {
 
       this.log.verbose(`Added attestation for slot ${slotNumber} from ${address}`);
     }
+
+    // TODO: set these to pending or something ????
+    this.metrics.recordAddedObjects(attestations.length);
     return Promise.resolve();
   }
 
+  #getNumberOfAttestationsInSlot(slot: bigint): number {
+    let total = 0;
+    const slotAttestationMap = getSlotOrDefault(this.attestations, slot);
+
+    if (slotAttestationMap) {
+      for (const proposalAttestationMap of slotAttestationMap.values() ?? []) {
+        total += proposalAttestationMap.size;
+      }
+    }
+    return total;
+  }
+
   public deleteAttestationsForSlot(slot: bigint): Promise<void> {
-    // TODO(md): check if this will free the memory of the inner hash map
+    // We count the number of attestations we are removing
+    const numberOfAttestations = this.#getNumberOfAttestationsInSlot(slot);
+
     this.attestations.delete(slot);
-    this.log.verbose(`Removed attestation for slot ${slot}`);
+    this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot}`);
+
+    this.metrics.recordRemovedObjects(numberOfAttestations);
     return Promise.resolve();
   }
 
   public deleteAttestationsForSlotAndProposal(slot: bigint, proposalId: string): Promise<void> {
-    const slotAttestationMap = this.attestations.get(slot);
+    const slotAttestationMap = getSlotOrDefault(this.attestations, slot);
     if (slotAttestationMap) {
-      slotAttestationMap.delete(proposalId);
-      this.log.verbose(`Removed attestation for slot ${slot}`);
+      if (slotAttestationMap.has(proposalId)) {
+        const numberOfAttestations = slotAttestationMap.get(proposalId)?.size ?? 0;
+
+        slotAttestationMap.delete(proposalId);
+
+        this.log.verbose(`Removed ${numberOfAttestations} attestations for slot ${slot} and proposal ${proposalId}`);
+        this.metrics.recordRemovedObjects(numberOfAttestations);
+      }
     }
     return Promise.resolve();
   }
@@ -68,6 +98,7 @@ export class InMemoryAttestationPool implements AttestationPool {
         }
       }
     }
+    this.metrics.recordRemovedObjects(attestations.length);
     return Promise.resolve();
   }
 }

yarn-project/p2p/src/mem_pools/epoch_proof_quote_pool/memory_epoch_proof_quote_pool.test.ts

@@ -1,19 +1,31 @@
-import { mockEpochProofQuote } from '@aztec/circuit-types';
+import { type EpochProofQuote, mockEpochProofQuote } from '@aztec/circuit-types';
+import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
+import { type MockProxy, mock } from 'jest-mock-extended';
+
+import { type PoolInstrumentation } from '../instrumentation.js';
 import { MemoryEpochProofQuotePool } from './memory_epoch_proof_quote_pool.js';
 
 describe('MemoryEpochProofQuotePool', () => {
   let pool: MemoryEpochProofQuotePool;
 
+  let metricsMock: MockProxy<PoolInstrumentation<EpochProofQuote>>;
+
   beforeEach(() => {
-    pool = new MemoryEpochProofQuotePool();
+    const telemetry = new NoopTelemetryClient();
+    pool = new MemoryEpochProofQuotePool(telemetry);
+
+    metricsMock = mock<PoolInstrumentation<EpochProofQuote>>();
+    (pool as any).metrics = metricsMock;
   });
 
   it('should add/get quotes to/from pool', () => {
     const quote = mockEpochProofQuote(5n);
 
     pool.addQuote(quote);
 
+    expect(metricsMock.recordAddedObjects).toHaveBeenCalledWith(1);
+
     const quotes = pool.getQuotes(quote.payload.epochToProve);
 
     expect(quotes).toHaveLength(1);
@@ -36,13 +48,16 @@ describe('MemoryEpochProofQuotePool', () => {
 
     const quotes3 = pool.getQuotes(3n);
     const quotesForEpoch3 = proofQuotes.filter(x => x.payload.epochToProve === 3n);
+    const quotesForEpoch2 = proofQuotes.filter(x => x.payload.epochToProve === 2n);
 
     expect(quotes3).toHaveLength(quotesForEpoch3.length);
     expect(quotes3).toEqual(quotesForEpoch3);
 
     // should delete all quotes for epochs 2 and 3
     pool.deleteQuotesToEpoch(3n);
 
+    expect(metricsMock.recordRemovedObjects).toHaveBeenCalledWith(quotesForEpoch2.length + quotesForEpoch3.length);
+
     expect(pool.getQuotes(2n)).toHaveLength(0);
     expect(pool.getQuotes(3n)).toHaveLength(0);