feat: add insturmentation to attestation and epoch quote mem pools (#…

…9055)

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

@@ -5,23 +5,27 @@ import { BlockAttestation } from './block_attestation.js';
 import { makeBlockAttestation } from './mocks.js';
 
 describe('Block Attestation serialization / deserialization', () => {
+  const checkEquivalence = (serialized: BlockAttestation, deserialized: BlockAttestation) => {
+    expect(deserialized.getSize()).toEqual(serialized.getSize());
+    expect(deserialized).toEqual(serialized);
+  };
+
   it('Should serialize / deserialize', () => {
     const attestation = makeBlockAttestation();
 
     const serialized = attestation.toBuffer();
     const deserialized = BlockAttestation.fromBuffer(serialized);
-
-    expect(deserialized).toEqual(attestation);
+    checkEquivalence(attestation, deserialized);
   });
 
   it('Should serialize / deserialize + recover sender', () => {
     const account = Secp256k1Signer.random();
 
-    const proposal = makeBlockAttestation(account);
-    const serialized = proposal.toBuffer();
+    const attestation = makeBlockAttestation(account);
+    const serialized = attestation.toBuffer();
     const deserialized = BlockAttestation.fromBuffer(serialized);
 
-    expect(deserialized).toEqual(proposal);
+    checkEquivalence(attestation, deserialized);
 
     // Recover signature
     const sender = deserialized.getSender();

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.test.ts

@@ -5,13 +5,17 @@ import { BlockProposal } from './block_proposal.js';
 import { makeBlockProposal } from './mocks.js';
 
 describe('Block Proposal serialization / deserialization', () => {
+  const checkEquivalence = (serialized: BlockProposal, deserialized: BlockProposal) => {
+    expect(deserialized.getSize()).toEqual(serialized.getSize());
+    expect(deserialized).toEqual(serialized);
+  };
+
   it('Should serialize / deserialize', () => {
     const proposal = makeBlockProposal();
 
     const serialized = proposal.toBuffer();
     const deserialized = BlockProposal.fromBuffer(serialized);
-
-    expect(deserialized).toEqual(proposal);
+    checkEquivalence(proposal, deserialized);
   });
 
   it('Should serialize / deserialize + recover sender', () => {
@@ -21,7 +25,7 @@ describe('Block Proposal serialization / deserialization', () => {
     const serialized = proposal.toBuffer();
     const deserialized = BlockProposal.fromBuffer(serialized);
 
-    expect(deserialized).toEqual(proposal);
+    checkEquivalence(proposal, deserialized);
 
     // Recover signature
     const sender = deserialized.getSender();

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,8 @@ import { TxHash } from '../tx/tx_hash.js';
 import { type Signable } from './signature_utils.js';
 
 export class ConsensusPayload implements Signable {
+  private size: number | undefined;
+
   constructor(
     /** The block header the attestation is made over */
     public readonly header: Header,
@@ -31,7 +33,9 @@ export class ConsensusPayload implements Signable {
   }
 
   toBuffer(): Buffer {
-    return serializeToBuffer([this.header, this.archive, this.txHashes.length, this.txHashes]);
+    const buffer = serializeToBuffer([this.header, this.archive, this.txHashes.length, this.txHashes]);
+    this.size = buffer.length;
+    return buffer;
   }
 
   static fromBuffer(buf: Buffer | BufferReader): ConsensusPayload {
@@ -50,4 +54,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.size = this.toBuffer().length;
+    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.test.ts

@@ -19,11 +19,18 @@ describe('epoch proof quote', () => {
     quote = new EpochProofQuote(payload, Signature.random());
   });
 
+  const checkEquivalence = (serialized: EpochProofQuote, deserialized: EpochProofQuote) => {
+    expect(deserialized.getSize()).toEqual(serialized.getSize());
+    expect(deserialized).toEqual(serialized);
+  };
+
   it('should serialize and deserialize from buffer', () => {
-    expect(EpochProofQuote.fromBuffer(quote.toBuffer())).toEqual(quote);
+    const deserialised = EpochProofQuote.fromBuffer(quote.toBuffer());
+    checkEquivalence(quote, deserialised);
   });
 
   it('should serialize and deserialize from JSON', () => {
-    expect(EpochProofQuote.fromJSON(quote.toJSON())).toEqual(quote);
+    const deserialised = EpochProofQuote.fromJSON(quote.toJSON());
+    checkEquivalence(quote, deserialised);
   });
 });

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.test.ts

@@ -20,16 +20,21 @@ describe('eth signature', () => {
     signature = signer.sign(message);
   });
 
+  const checkEquivalence = (serialized: Signature, deserialized: Signature) => {
+    expect(deserialized.getSize()).toEqual(serialized.getSize());
+    expect(deserialized).toEqual(serialized);
+  };
+
   it('should serialize / deserialize to buffer', () => {
     const serialized = signature.toBuffer();
     const deserialized = Signature.fromBuffer(serialized);
-    expect(deserialized).toEqual(signature);
+    checkEquivalence(signature, deserialized);
   });
 
   it('should serialize / deserialize real signature to hex string', () => {
     const serialized = signature.to0xString();
     const deserialized = Signature.from0xString(serialized);
-    expect(deserialized).toEqual(signature);
+    checkEquivalence(signature, deserialized);
   });
 
   it('should recover signer from signature', () => {
@@ -41,13 +46,13 @@ describe('eth signature', () => {
     const signature = new Signature(Buffer32.random(), Buffer32.random(), 1, false);
     const serialized = signature.to0xString();
     const deserialized = Signature.from0xString(serialized);
-    expect(deserialized).toEqual(signature);
+    checkEquivalence(signature, deserialized);
   });
 
   it('should serialize / deserialize to hex string with 2-digit v', () => {
     const signature = new Signature(Buffer32.random(), Buffer32.random(), 26, false);
     const serialized = signature.to0xString();
     const deserialized = Signature.from0xString(serialized);
-    expect(deserialized).toEqual(signature);
+    checkEquivalence(signature, deserialized);
   });
 });

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

@@ -19,6 +19,9 @@ export type ViemSignature = {
  * Contains a signature split into it's primary components (r,s,v)
  */
 export class Signature {
+  // Cached values
+  private size: number | undefined;
+
   constructor(
     /** The r value of the signature */
     public readonly r: Buffer32,
@@ -73,7 +76,19 @@ export class Signature {
   }
 
   toBuffer(): Buffer {
-    return serializeToBuffer([this.r, this.s, this.v]);
+    const buffer = serializeToBuffer([this.r, this.s, this.v]);
+    this.size = buffer.length;
+    return buffer;
+  }
+
+  getSize(): number {
+    // We cache size to avoid recalculating it
+    if (this.size) {
+      return this.size;
+    }
+
+    this.size = this.toBuffer().length;
+    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);
   });