Skip to content

Commit

Permalink
feat: add insturmentation to attestation and epoch quote mem pools (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
Maddiaa0 authored Oct 9, 2024
1 parent 7677ca5 commit 7dfa295
Show file tree
Hide file tree
Showing 20 changed files with 301 additions and 120 deletions.
14 changes: 9 additions & 5 deletions yarn-project/circuit-types/src/p2p/block_attestation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
4 changes: 4 additions & 0 deletions yarn-project/circuit-types/src/p2p/block_attestation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
10 changes: 7 additions & 3 deletions yarn-project/circuit-types/src/p2p/block_proposal.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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();
Expand Down
4 changes: 4 additions & 0 deletions yarn-project/circuit-types/src/p2p/block_proposal.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
19 changes: 18 additions & 1 deletion yarn-project/circuit-types/src/p2p/consensus_payload.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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 {
Expand All @@ -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;
}
}
7 changes: 7 additions & 0 deletions yarn-project/circuit-types/src/p2p/gossipable.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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 {
Expand Down Expand Up @@ -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} }`;
}
Expand Down
13 changes: 9 additions & 4 deletions yarn-project/foundation/src/eth-signature/eth_signature.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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);
});
});
17 changes: 16 additions & 1 deletion yarn-project/foundation/src/eth-signature/eth_signature.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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}` {
Expand Down
4 changes: 2 additions & 2 deletions yarn-project/p2p/src/client/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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 () => {
Expand All @@ -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);
Expand All @@ -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);
});
Expand Down Expand Up @@ -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);
});
Expand Down Expand Up @@ -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);
});
Expand Down
Loading

0 comments on commit 7dfa295

Please sign in to comment.