fix: attestations are requested based on their proposal not just slot (…

…#8442)
AztecProtocol · Sep 27, 2024 · 1d42b8a · 1d42b8a
1 parent 3302c48
commit 1d42b8a
Show file tree

Hide file tree

Showing 6 changed files with 114 additions and 26 deletions.
diff --git 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<BlockAttestation[]>;
+  getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]>;
 }
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
@@ -1,3 +1,5 @@
+import { Fr } from '@aztec/foundation/fields';
+
 import { type PrivateKeyAccount } from 'viem';
 
 import { InMemoryAttestationPool } from './memory_attestation_pool.js';
@@ -16,19 +18,22 @@ 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);
 
     // 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
@@ -4,30 +4,34 @@ import { createDebugLogger } from '@aztec/foundation/log';
 import { type AttestationPool } from './attestation_pool.js';
 
 export class InMemoryAttestationPool implements AttestationPool {
-  private attestations: Map</*slot=*/ bigint, Map</*address=*/ string, BlockAttestation>>;
+  private attestations: Map</*slot=*/ bigint, Map</*proposalId*/ string, Map</*address=*/ string, BlockAttestation>>>;
 
   constructor(private log = createDebugLogger('aztec:attestation_pool')) {
     this.attestations = new Map();
   }
 
-  public getAttestationsForSlot(slot: bigint): Promise<BlockAttestation[]> {
+  public getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
     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<void> {
     for (const attestation of attestations) {
       // 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<void> {
+    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<void> {
     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<bigint, Map<string, BlockAttestation>>,
+  map: Map<bigint, Map<string, Map<string, BlockAttestation>>>,
   slot: bigint,
-): Map<string, BlockAttestation> {
+): Map<string, Map<string, BlockAttestation>> {
   if (!map.has(slot)) {
-    map.set(slot, new Map<string, BlockAttestation>());
+    map.set(slot, new Map<string, Map<string, BlockAttestation>>());
   }
   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<string, Map<string, BlockAttestation>>,
+  proposalId: string,
+): Map<string, BlockAttestation> {
+  if (!map.has(proposalId)) {
+    map.set(proposalId, new Map<string, BlockAttestation>());
+  }
+  return map.get(proposalId)!;
+}
diff --git 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<BlockAttestation> => {
+export const mockAttestation = async (
+  signer: PrivateKeyAccount,
+  slot: number = 0,
+  archive: Fr = Fr.random(),
+): Promise<BlockAttestation> => {
   // 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
@@ -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<BlockAttestation[]>;
+  getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]>;
 
   /**
    * 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<BlockAttestation[]> {
-    return Promise.resolve(this.attestationPool.getAttestationsForSlot(slot));
+  public getAttestationsForSlot(slot: bigint, proposalId: string): Promise<BlockAttestation[]> {
+    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
@@ -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}`);