diff --git a/yarn-project/p2p/src/attestation_pool/attestation_pool.ts b/yarn-project/p2p/src/attestation_pool/attestation_pool.ts index a4fcb993c6e..cdfe8729911 100644 --- a/yarn-project/p2p/src/attestation_pool/attestation_pool.ts +++ b/yarn-project/p2p/src/attestation_pool/attestation_pool.ts @@ -36,7 +36,8 @@ export interface AttestationPool { * Retrieve all of the attestations observed pertaining to a given slot * * @param slot - The slot to query + * @param proposalId - The proposal to query * @return BlockAttestations */ - getAttestationsForSlot(slot: bigint): Promise; + getAttestationsForSlot(slot: bigint, proposalId: string): Promise; } diff --git a/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.test.ts b/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.test.ts index 98c51e2dd0f..e05d068940b 100644 --- a/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.test.ts +++ b/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.test.ts @@ -1,3 +1,5 @@ +import { Fr } from '@aztec/foundation/fields'; + import { type PrivateKeyAccount } from 'viem'; import { InMemoryAttestationPool } from './memory_attestation_pool.js'; @@ -16,11 +18,14 @@ describe('MemoryAttestationPool', () => { it('should add attestation to pool', async () => { const slotNumber = 420; - const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber))); + const archive = Fr.random(); + const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber, archive))); + + const proposalId = attestations[0].p2pMessageIdentifier.toString(); await ap.addAttestations(attestations); - const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber)); + const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); expect(retreivedAttestations.length).toBe(NUMBER_OF_SIGNERS_PER_TEST); expect(retreivedAttestations).toEqual(attestations); @@ -28,7 +33,7 @@ describe('MemoryAttestationPool', () => { // Delete by slot await ap.deleteAttestationsForSlot(BigInt(slotNumber)); - const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber)); + const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); expect(retreivedAttestationsAfterDelete.length).toBe(0); }); @@ -40,8 +45,9 @@ describe('MemoryAttestationPool', () => { for (const attestation of attestations) { const slot = attestation.payload.header.globalVariables.slotNumber; + const proposalId = attestation.p2pMessageIdentifier.toString(); - const retreivedAttestations = await ap.getAttestationsForSlot(slot.toBigInt()); + const retreivedAttestations = await ap.getAttestationsForSlot(slot.toBigInt(), proposalId); expect(retreivedAttestations.length).toBe(1); expect(retreivedAttestations[0]).toEqual(attestation); expect(retreivedAttestations[0].payload.header.globalVariables.slotNumber).toEqual(slot); @@ -50,17 +56,55 @@ describe('MemoryAttestationPool', () => { it('Should delete attestations', async () => { const slotNumber = 420; - const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber))); + const archive = Fr.random(); + const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber, archive))); + const proposalId = attestations[0].p2pMessageIdentifier.toString(); await ap.addAttestations(attestations); - const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber)); + 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); - const gottenAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber)); + const gottenAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); expect(gottenAfterDelete.length).toBe(0); }); + + it('Should blanket delete attestations per slot', async () => { + const slotNumber = 420; + const archive = Fr.random(); + const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber, archive))); + const proposalId = attestations[0].p2pMessageIdentifier.toString(); + + await ap.addAttestations(attestations); + + const retreivedAttestations = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); + expect(retreivedAttestations.length).toBe(NUMBER_OF_SIGNERS_PER_TEST); + expect(retreivedAttestations).toEqual(attestations); + + await ap.deleteAttestationsForSlot(BigInt(slotNumber)); + + const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); + expect(retreivedAttestationsAfterDelete.length).toBe(0); + }); + + it('Should blanket delete attestations per slot and proposal', async () => { + const slotNumber = 420; + const archive = Fr.random(); + const attestations = await Promise.all(signers.map(signer => mockAttestation(signer, slotNumber, archive))); + const proposalId = attestations[0].p2pMessageIdentifier.toString(); + + await ap.addAttestations(attestations); + + 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); + + const retreivedAttestationsAfterDelete = await ap.getAttestationsForSlot(BigInt(slotNumber), proposalId); + expect(retreivedAttestationsAfterDelete.length).toBe(0); + }); }); diff --git a/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.ts b/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.ts index c71a638d99c..f6c35a0b485 100644 --- a/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.ts +++ b/yarn-project/p2p/src/attestation_pool/memory_attestation_pool.ts @@ -4,19 +4,21 @@ import { createDebugLogger } from '@aztec/foundation/log'; import { type AttestationPool } from './attestation_pool.js'; export class InMemoryAttestationPool implements AttestationPool { - private attestations: Map>; + private attestations: Map>>; constructor(private log = createDebugLogger('aztec:attestation_pool')) { this.attestations = new Map(); } - public getAttestationsForSlot(slot: bigint): Promise { + public getAttestationsForSlot(slot: bigint, proposalId: string): Promise { const slotAttestationMap = this.attestations.get(slot); if (slotAttestationMap) { - return Promise.resolve(Array.from(slotAttestationMap.values())); - } else { - return Promise.resolve([]); + const proposalAttestationMap = slotAttestationMap.get(proposalId); + if (proposalAttestationMap) { + return Promise.resolve(Array.from(proposalAttestationMap.values())); + } } + return Promise.resolve([]); } public addAttestations(attestations: BlockAttestation[]): Promise { @@ -24,10 +26,12 @@ export class InMemoryAttestationPool implements AttestationPool { // Perf: order and group by slot before insertion const slotNumber = attestation.payload.header.globalVariables.slotNumber; + const proposalId = attestation.p2pMessageIdentifier.toString(); const address = attestation.getSender(); const slotAttestationMap = getSlotOrDefault(this.attestations, slotNumber.toBigInt()); - slotAttestationMap.set(address.toString(), attestation); + const proposalAttestationMap = getProposalOrDefault(slotAttestationMap, proposalId); + proposalAttestationMap.set(address.toString(), attestation); this.log.verbose(`Added attestation for slot ${slotNumber} from ${address}`); } @@ -41,14 +45,27 @@ export class InMemoryAttestationPool implements AttestationPool { return Promise.resolve(); } + public deleteAttestationsForSlotAndProposal(slot: bigint, proposalId: string): Promise { + const slotAttestationMap = this.attestations.get(slot); + if (slotAttestationMap) { + slotAttestationMap.delete(proposalId); + this.log.verbose(`Removed attestation for slot ${slot}`); + } + return Promise.resolve(); + } + public deleteAttestations(attestations: BlockAttestation[]): Promise { for (const attestation of attestations) { const slotNumber = attestation.payload.header.globalVariables.slotNumber; const slotAttestationMap = this.attestations.get(slotNumber.toBigInt()); if (slotAttestationMap) { - const address = attestation.getSender(); - slotAttestationMap.delete(address.toString()); - this.log.verbose(`Deleted attestation for slot ${slotNumber} from ${address}`); + const proposalId = attestation.p2pMessageIdentifier.toString(); + const proposalAttestationMap = getProposalOrDefault(slotAttestationMap, proposalId); + if (proposalAttestationMap) { + const address = attestation.getSender(); + proposalAttestationMap.delete(address.toString()); + this.log.debug(`Deleted attestation for slot ${slotNumber} from ${address}`); + } } } return Promise.resolve(); @@ -59,13 +76,34 @@ export class InMemoryAttestationPool implements AttestationPool { * Get Slot or Default * * Fetch the slot mapping, if it does not exist, then create a mapping and return it + * @param map - The map to fetch from + * @param slot - The slot to fetch + * @returns The slot mapping */ function getSlotOrDefault( - map: Map>, + map: Map>>, slot: bigint, -): Map { +): Map> { if (!map.has(slot)) { - map.set(slot, new Map()); + map.set(slot, new Map>()); } return map.get(slot)!; } + +/** + * Get Proposal or Default + * + * Fetch the proposal mapping, if it does not exist, then create a mapping and return it + * @param map - The map to fetch from + * @param proposalId - The proposal id to fetch + * @returns The proposal mapping + */ +function getProposalOrDefault( + map: Map>, + proposalId: string, +): Map { + if (!map.has(proposalId)) { + map.set(proposalId, new Map()); + } + return map.get(proposalId)!; +} diff --git a/yarn-project/p2p/src/attestation_pool/mocks.ts b/yarn-project/p2p/src/attestation_pool/mocks.ts index a0c3e1d05f3..bf75e7c3905 100644 --- a/yarn-project/p2p/src/attestation_pool/mocks.ts +++ b/yarn-project/p2p/src/attestation_pool/mocks.ts @@ -22,10 +22,13 @@ export const generateAccount = () => { * @param slot The slot number the attestation is for * @returns A Block Attestation */ -export const mockAttestation = async (signer: PrivateKeyAccount, slot: number = 0): Promise => { +export const mockAttestation = async ( + signer: PrivateKeyAccount, + slot: number = 0, + archive: Fr = Fr.random(), +): Promise => { // Use arbitrary numbers for all other than slot const header = makeHeader(1, 2, slot); - const archive = Fr.random(); const txs = [0, 1, 2, 3, 4, 5].map(() => TxHash.random()); const payload = new ConsensusPayload(header, archive, txs); diff --git a/yarn-project/p2p/src/client/p2p_client.ts b/yarn-project/p2p/src/client/p2p_client.ts index 63c5b453f0a..4738d153006 100644 --- a/yarn-project/p2p/src/client/p2p_client.ts +++ b/yarn-project/p2p/src/client/p2p_client.ts @@ -58,9 +58,10 @@ export interface P2P { * Queries the Attestation pool for attestations for the given slot * * @param slot - the slot to query + * @param proposalId - the proposal id to query * @returns BlockAttestations */ - getAttestationsForSlot(slot: bigint): Promise; + getAttestationsForSlot(slot: bigint, proposalId: string): Promise; /** * Registers a callback from the validator client that determines how to behave when @@ -281,8 +282,8 @@ export class P2PClient implements P2P { return this.p2pService.propagate(proposal); } - public getAttestationsForSlot(slot: bigint): Promise { - return Promise.resolve(this.attestationPool.getAttestationsForSlot(slot)); + public getAttestationsForSlot(slot: bigint, proposalId: string): Promise { + return Promise.resolve(this.attestationPool.getAttestationsForSlot(slot, proposalId)); } // REVIEW: https://github.com/AztecProtocol/aztec-packages/issues/7963 diff --git a/yarn-project/validator-client/src/validator.ts b/yarn-project/validator-client/src/validator.ts index a32a7d4f9e7..49c24b1e211 100644 --- a/yarn-project/validator-client/src/validator.ts +++ b/yarn-project/validator-client/src/validator.ts @@ -141,12 +141,13 @@ export class ValidatorClient implements Validator { const slot = proposal.payload.header.globalVariables.slotNumber.toBigInt(); this.log.info(`Waiting for ${numberOfRequiredAttestations} attestations for slot: ${slot}`); + const proposalId = proposal.p2pMessageIdentifier.toString(); const myAttestation = await this.validationService.attestToProposal(proposal); const startTime = Date.now(); while (true) { - const attestations = [myAttestation, ...(await this.p2pClient.getAttestationsForSlot(slot))]; + const attestations = [myAttestation, ...(await this.p2pClient.getAttestationsForSlot(slot, proposalId))]; if (attestations.length >= numberOfRequiredAttestations) { this.log.info(`Collected all ${numberOfRequiredAttestations} attestations for slot, ${slot}`);