refactor(pxe): append only NoteStore and PrivateEventStore (#23785)

Refactors mainly NoteStore so that discovered nullifiers are modeled as
append-only rows instead of mutating stored notes. Note the store isn't
100% append-only: adding scopes to a note still causes a modification. I
deemed this less cumbersome than introducing more indexes just to track
scopes.

Includes some minor touches to PrivateEventStore mostly to align style,
since it was already mostly in the shape we need it to.

Context: this is preliminary work for the "time-travelling pxe db"
project. We need reorg semantics to be consistent between stores,
dealing with mutation on reorgs makes the whole model harder to reason
about.

You'll see a longer than expected commit history: this is because I
originally set out to change the way we handle reorgs much more
fundamentally, by making the stores tolerate multiple versions of notes
and events as long as they are recorded as having been originated from
different blocks (which could happen in reorg scenarios). Implementing
said approach made me realize we don't win much from the additional
complexity: in order to filter out "orphaned" data and eventually being
able to reap it, we would need to keep track of height->hash mappings
from genesis.

At that point, the delete-on-chain-prune original approach seemed more
sensible, which ostensibly reduced the scope of this change to what I
describe above.

Author: mverzilli

yarn-project/pxe/src/block_synchronizer/block_synchronizer.test.ts

@@ -4,18 +4,24 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import { openTmpStore } from '@aztec/kv-store/lmdb-v2';
 import { L2TipsKVStore } from '@aztec/kv-store/stores';
+import { EventSelector } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import {
   type BlockData,
   BlockHash,
   GENESIS_BLOCK_HEADER_HASH,
   GENESIS_CHECKPOINT_HEADER_HASH,
   L2Block,
+  type L2BlockId,
   type L2BlockStream,
+  makeL2BlockId,
+  makeL2CheckpointId,
 } from '@aztec/stdlib/block';
 import { Checkpoint, L1PublishedData, PublishedCheckpoint } from '@aztec/stdlib/checkpoint';
-import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { AztecNode, BlockResponse } from '@aztec/stdlib/interfaces/client';
+import { NoteDao, NoteStatus } from '@aztec/stdlib/note';
+import { TxHash } from '@aztec/stdlib/tx';
 
-import { jest } from '@jest/globals';
 import { type MockProxy, mock } from 'jest-mock-extended';
 
 import type { BlockSynchronizerConfig } from '../config/index.js';
@@ -25,6 +31,12 @@ import { NoteStore } from '../storage/note_store/note_store.js';
 import { PrivateEventStore } from '../storage/private_event_store/private_event_store.js';
 import { BlockSynchronizer } from './block_synchronizer.js';
 
+// `AztecNode.getBlock` is generic over its include-options; `Parameters`/`ReturnType` collapse that
+// type parameter to its `BlockIncludeOptions` constraint, yielding a concrete signature the mock can satisfy.
+type NodeGetBlockMock = jest.MockedFunction<
+  (...args: Parameters<AztecNode['getBlock']>) => ReturnType<AztecNode['getBlock']>
+>;
+
 describe('BlockSynchronizer', () => {
   let synchronizer: BlockSynchronizer;
   let store: AztecAsyncKVStore;
@@ -33,6 +45,7 @@ describe('BlockSynchronizer', () => {
   let noteStore: NoteStore;
   let privateEventStore: PrivateEventStore;
   let aztecNode: MockProxy<AztecNode>;
+  let getBlock: NodeGetBlockMock;
   let blockStream: MockProxy<L2BlockStream>;
   let contractSyncService: MockProxy<ContractSyncService>;
 
@@ -55,10 +68,44 @@ describe('BlockSynchronizer', () => {
     );
   };
 
+  // Builds the BlockResponse the node returns for a block (handlers read its header to update the anchor).
+  const blockResponse = async (block: L2Block): Promise<BlockResponse> => ({
+    header: block.header,
+    archive: block.archive,
+    hash: await block.hash(),
+    checkpointNumber: block.checkpointNumber,
+    indexWithinCheckpoint: block.indexWithinCheckpoint,
+    number: block.number,
+  });
+
+  // Builds a note DAO anchored to the given block id (caller adds it to the store and commits).
+  const noteAt = (contract: AztecAddress, block: L2BlockId): Promise<NoteDao> =>
+    NoteDao.random({ contractAddress: contract, l2BlockNumber: block.number, l2BlockHash: block.hash });
+
+  // Stores one private event anchored to the given block id under the 'event-job' (caller commits).
+  const storeEvent = (contract: AztecAddress, scope: AztecAddress, eventId: Fr, block: L2BlockId) =>
+    privateEventStore.storePrivateEventLog(
+      EventSelector.random(),
+      Fr.random(),
+      [Fr.random()],
+      eventId,
+      {
+        contractAddress: contract,
+        scope,
+        txHash: TxHash.random(),
+        l2BlockNumber: block.number,
+        l2BlockHash: BlockHash.fromString(block.hash),
+        txIndexInBlock: 0,
+        eventIndexInTx: 0,
+      },
+      'event-job',
+    );
+
   beforeEach(async () => {
     store = await openTmpStore('test');
     blockStream = mock<L2BlockStream>();
     aztecNode = mock<AztecNode>();
+    getBlock = aztecNode.getBlock as NodeGetBlockMock;
     tipsStore = new L2TipsKVStore(store, 'pxe', GENESIS_BLOCK_HEADER_HASH);
     anchorBlockStore = new AnchorBlockStore(store);
     noteStore = new NoteStore(store);
@@ -75,66 +122,26 @@ describe('BlockSynchronizer', () => {
     expect(obtainedHeader.equals(block.header)).toBe(true);
   });
 
-  it('removes notes from db on a reorg', async () => {
-    const rollback = jest.spyOn(noteStore, 'rollback').mockImplementation(() => Promise.resolve());
-    const block3Hash = Fr.fromString('0x3');
-    aztecNode.getBlock.mockImplementation(async (block: any) => {
-      if (block instanceof BlockHash && block.equals(block3Hash)) {
-        const b = await L2Block.random(BlockNumber(3));
-        return {
-          header: b.header,
-          archive: b.archive,
-          hash: await b.hash(),
-          checkpointNumber: b.checkpointNumber,
-          indexWithinCheckpoint: b.indexWithinCheckpoint,
-          number: b.number,
-        } as any;
-      }
-      return undefined;
-    });
-
-    await synchronizer.handleBlockStreamEvent({
-      type: 'blocks-added',
-      blocks: await timesParallel(5, i => L2Block.random(BlockNumber(i))),
-    });
-    await synchronizer.handleBlockStreamEvent({
-      type: 'chain-pruned',
-      block: { number: BlockNumber(3), hash: block3Hash.toString() },
-      checkpoint: { number: CheckpointNumber.ZERO, hash: GENESIS_CHECKPOINT_HEADER_HASH.toString() },
-    });
-
-    expect(rollback).toHaveBeenCalledWith(3, 4);
-  });
-
-  it('removes private events from db on a reorg', async () => {
-    const rollback = jest.spyOn(privateEventStore, 'rollback').mockImplementation(() => Promise.resolve());
-    const block3Hash = Fr.fromString('0x3');
-    aztecNode.getBlock.mockImplementation(async (block: any) => {
-      if (block instanceof BlockHash && block.equals(block3Hash)) {
-        const b = await L2Block.random(BlockNumber(3));
-        return {
-          header: b.header,
-          archive: b.archive,
-          hash: await b.hash(),
-          checkpointNumber: b.checkpointNumber,
-          indexWithinCheckpoint: b.indexWithinCheckpoint,
-          number: b.number,
-        } as any;
-      }
-      return undefined;
-    });
+  it('updates anchor block on a reorg', async () => {
+    const reorgBlock = await L2Block.random(BlockNumber(3));
+    const reorgResponse = await blockResponse(reorgBlock);
+    getBlock.mockImplementation(param =>
+      Promise.resolve(param instanceof BlockHash && param.equals(reorgResponse.hash) ? reorgResponse : undefined),
+    );
 
     await synchronizer.handleBlockStreamEvent({
       type: 'blocks-added',
       blocks: await timesParallel(5, i => L2Block.random(BlockNumber(i))),
     });
     await synchronizer.handleBlockStreamEvent({
       type: 'chain-pruned',
-      block: { number: BlockNumber(3), hash: block3Hash.toString() },
-      checkpoint: { number: CheckpointNumber.ZERO, hash: GENESIS_CHECKPOINT_HEADER_HASH.toString() },
+      block: makeL2BlockId(reorgBlock.number, reorgResponse.hash.toString()),
+      checkpoint: makeL2CheckpointId(CheckpointNumber.ZERO, GENESIS_CHECKPOINT_HEADER_HASH.toString()),
     });
 
-    expect(rollback).toHaveBeenCalledWith(3, 4);
+    // The anchor block should be updated to the reorg block header.
+    const obtainedHeader = await anchorBlockStore.getBlockHeader();
+    expect(obtainedHeader.equals(reorgBlock.header)).toBe(true);
   });
 
   describe('stop', () => {
@@ -182,6 +189,139 @@ describe('BlockSynchronizer', () => {
     });
   });
 
+  describe('delete-on-prune', () => {
+    it('chain-pruned deletes rows anchored above the fork and keeps rows at or below it', async () => {
+      const contract = await AztecAddress.random();
+      const scope = await AztecAddress.random();
+
+      // Block 3 is the fork point (a real block the node still serves); 4 and 5 are on the abandoned fork.
+      const forkBlock = await L2Block.random(BlockNumber(3));
+      const block3 = makeL2BlockId(forkBlock.number, (await forkBlock.hash()).toString());
+      const block4 = makeL2BlockId(BlockNumber(4), Fr.random().toString());
+      const block5 = makeL2BlockId(BlockNumber(5), Fr.random().toString());
+
+      // Seed a note at each block, anchored to that block's id.
+      const noteAt3 = await noteAt(contract, block3);
+      const noteAt4 = await noteAt(contract, block4);
+      const noteAt5 = await noteAt(contract, block5);
+      await noteStore.addNotes([noteAt3, noteAt4, noteAt5], scope, 'note-job');
+      await noteStore.commit('note-job');
+
+      // Seed an event at each block.
+      const eventIdAt3 = Fr.random();
+      const eventIdAt4 = Fr.random();
+      const eventIdAt5 = Fr.random();
+      await storeEvent(contract, scope, eventIdAt3, block3);
+      await storeEvent(contract, scope, eventIdAt4, block4);
+      await storeEvent(contract, scope, eventIdAt5, block5);
+      await privateEventStore.commit('event-job');
+
+      // Set the anchor to block 5 so the prune guard passes.
+      const anchorBlock5 = await L2Block.random(BlockNumber(5));
+      await anchorBlockStore.setHeader(anchorBlock5.header);
+
+      // The node serves the fork-point block; it becomes the new anchor after the prune.
+      const forkResponse = await blockResponse(forkBlock);
+      getBlock.mockImplementation(param =>
+        Promise.resolve(param instanceof BlockHash && param.equals(forkResponse.hash) ? forkResponse : undefined),
+      );
+
+      // Prune back to block 3 (orphaning blocks 4 and 5).
+      await synchronizer.handleBlockStreamEvent({
+        type: 'chain-pruned',
+        block: block3,
+        checkpoint: makeL2CheckpointId(CheckpointNumber.ZERO, GENESIS_CHECKPOINT_HEADER_HASH.toString()),
+      });
+
+      // Rows at blocks 4 and 5 must be gone.
+      expect(await noteStore.nullifiersOfNotesAtBlock(4)).toHaveLength(0);
+      expect(await noteStore.nullifiersOfNotesAtBlock(5)).toHaveLength(0);
+      expect(await privateEventStore.eventIdsAtBlock(4)).toHaveLength(0);
+      expect(await privateEventStore.eventIdsAtBlock(5)).toHaveLength(0);
+
+      // Rows at block 3 (the fork point, not an orphan) must survive.
+      expect(await noteStore.nullifiersOfNotesAtBlock(3)).toEqual([noteAt3.siloedNullifier.toString()]);
+      expect(await privateEventStore.eventIdsAtBlock(3)).toEqual([eventIdAt3.toString()]);
+    });
+
+    it('chain-finalized does not delete any rows', async () => {
+      const contract = await AztecAddress.random();
+      const scope = await AztecAddress.random();
+
+      // Canonical rows at two heights: one below the finalized block, one at it.
+      const block8 = makeL2BlockId(BlockNumber(8), Fr.random().toString());
+      const block9 = makeL2BlockId(BlockNumber(9), Fr.random().toString());
+      const note8 = await noteAt(contract, block8);
+      const note9 = await noteAt(contract, block9);
+      await noteStore.addNotes([note8, note9], scope, 'note-job');
+      await noteStore.commit('note-job');
+
+      const eventId8 = Fr.random();
+      const eventId9 = Fr.random();
+      await storeEvent(contract, scope, eventId8, block8);
+      await storeEvent(contract, scope, eventId9, block9);
+      await privateEventStore.commit('event-job');
+
+      await synchronizer.handleBlockStreamEvent({
+        type: 'chain-finalized',
+        block: block9,
+      });
+
+      // Finalization is a no-op for storage under delete-on-prune, every row at and below the tip survives.
+      expect(await noteStore.nullifiersOfNotesAtBlock(8)).toEqual([note8.siloedNullifier.toString()]);
+      expect(await noteStore.nullifiersOfNotesAtBlock(9)).toEqual([note9.siloedNullifier.toString()]);
+      expect(await privateEventStore.eventIdsAtBlock(8)).toEqual([eventId8.toString()]);
+      expect(await privateEventStore.eventIdsAtBlock(9)).toEqual([eventId9.toString()]);
+    });
+
+    it('notes below the fork survive and remain queryable after a prune', async () => {
+      const contract = await AztecAddress.random();
+      const scope = await AztecAddress.random();
+
+      // Block 1 is the fork point (a real block the node still serves); 2 and 3 are on the abandoned fork.
+      const forkBlock = await L2Block.random(BlockNumber(1));
+      const block1 = makeL2BlockId(forkBlock.number, (await forkBlock.hash()).toString());
+      const block2 = makeL2BlockId(BlockNumber(2), Fr.random().toString());
+      const block3 = makeL2BlockId(BlockNumber(3), Fr.random().toString());
+
+      const noteAt1 = await noteAt(contract, block1);
+      const noteAt2 = await noteAt(contract, block2);
+      const noteAt3 = await noteAt(contract, block3);
+      await noteStore.addNotes([noteAt1, noteAt2, noteAt3], scope, 'note-job');
+      await noteStore.commit('note-job');
+
+      // Anchor at block 3.
+      const anchorBlock3 = await L2Block.random(BlockNumber(3));
+      await anchorBlockStore.setHeader(anchorBlock3.header);
+
+      // The node serves the fork-point block; it becomes the new anchor after the prune.
+      const forkResponse = await blockResponse(forkBlock);
+      getBlock.mockImplementation(param =>
+        Promise.resolve(param instanceof BlockHash && param.equals(forkResponse.hash) ? forkResponse : undefined),
+      );
+
+      // Prune back to block 1 (orphaning blocks 2 and 3).
+      await synchronizer.handleBlockStreamEvent({
+        type: 'chain-pruned',
+        block: block1,
+        checkpoint: makeL2CheckpointId(CheckpointNumber.ZERO, GENESIS_CHECKPOINT_HEADER_HASH.toString()),
+      });
+
+      // Blocks 2 and 3 deleted.
+      expect(await noteStore.nullifiersOfNotesAtBlock(2)).toHaveLength(0);
+      expect(await noteStore.nullifiersOfNotesAtBlock(3)).toHaveLength(0);
+
+      // Block 1 note still present and visible via getNotes.
+      expect(await noteStore.nullifiersOfNotesAtBlock(1)).toEqual([noteAt1.siloedNullifier.toString()]);
+      const found = await noteStore.getNotes(
+        { contractAddress: contract, scopes: [scope], status: NoteStatus.ACTIVE },
+        'read-job',
+      );
+      expect(found).toHaveLength(1);
+      expect(found[0].siloedNullifier.equals(noteAt1.siloedNullifier)).toBe(true);
+    });
+  });
+
   describe('syncChainTip config', () => {
     it('updates anchor on blocks-added when syncChainTip is proposed (default)', async () => {
       synchronizer = createSynchronizer({ syncChainTip: 'proposed' });
@@ -263,7 +403,7 @@ describe('BlockSynchronizer', () => {
 
       // Mock node to return block
       const provenBlock = await L2Block.random(BlockNumber(5));
-      aztecNode.getBlock.mockResolvedValue({ header: provenBlock.header } as any);
+      getBlock.mockResolvedValue(await blockResponse(provenBlock));
 
       await synchronizer.handleBlockStreamEvent({
         type: 'chain-proven',
@@ -283,7 +423,7 @@ describe('BlockSynchronizer', () => {
 
       // Mock node to return block
       const finalizedBlock = await L2Block.random(BlockNumber(10));
-      aztecNode.getBlock.mockResolvedValue({ header: finalizedBlock.header } as any);
+      getBlock.mockResolvedValue(await blockResponse(finalizedBlock));
 
       await synchronizer.handleBlockStreamEvent({
         type: 'chain-finalized',
@@ -301,17 +441,13 @@ describe('BlockSynchronizer', () => {
       const anchorBlock = await L2Block.random(BlockNumber(2));
       await anchorBlockStore.setHeader(anchorBlock.header);
 
-      const rollback = jest.spyOn(noteStore, 'rollback').mockImplementation(() => Promise.resolve());
-
       // Prune to block 3 (above anchor) - should be ignored
       await synchronizer.handleBlockStreamEvent({
         type: 'chain-pruned',
         block: { number: BlockNumber(3), hash: '0x3' },
-        checkpoint: { number: CheckpointNumber.ZERO, hash: GENESIS_CHECKPOINT_HEADER_HASH.toString() },
+        checkpoint: makeL2CheckpointId(CheckpointNumber.ZERO, GENESIS_CHECKPOINT_HEADER_HASH.toString()),
       });
 
-      expect(rollback).not.toHaveBeenCalled();
-
       // Anchor should be unchanged
       const obtainedHeader = await anchorBlockStore.getBlockHeader();
       expect(obtainedHeader.equals(anchorBlock.header)).toBe(true);

yarn-project/pxe/src/block_synchronizer/block_synchronizer.ts

@@ -9,8 +9,8 @@ import type { BlockHeader } from '@aztec/stdlib/tx';
 
 import type { BlockSynchronizerConfig } from '../config/index.js';
 import type { ContractSyncService } from '../contract_sync/contract_sync_service.js';
-import type { AnchorBlockStore } from '../storage/anchor_block_store/anchor_block_store.js';
-import type { NoteStore } from '../storage/note_store/note_store.js';
+import type { AnchorBlockStore } from '../storage/anchor_block_store/index.js';
+import type { NoteStore } from '../storage/note_store/index.js';
 import type { PrivateEventStore } from '../storage/private_event_store/private_event_store.js';
 import { blockStreamSourceFromAztecNode } from './block_stream_source.js';
 
@@ -113,7 +113,7 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
           return;
         }
 
-        this.log.warn(`Pruning data after block ${event.block.number} due to reorg`);
+        this.log.warn(`Pruning data after block ${event.block.number} due to reorg`, { pruneBlock: event.block });
 
         // Note that the following is not necessarily the anchor block that will be used in the transaction - if
         // the chain has already moved past the reorg, we'll also see blocks-added events that will push the anchor
@@ -129,8 +129,8 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
 
         // Operations are wrapped in a single transaction to ensure atomicity.
         await this.store.transactionAsync(async () => {
-          await this.noteStore.rollback(event.block.number, currentAnchorBlockNumber);
-          await this.privateEventStore.rollback(event.block.number, currentAnchorBlockNumber);
+          await this.noteStore.rollback(event.block.number);
+          await this.privateEventStore.rollback(event.block.number);
           await this.updateAnchorBlockHeader(newAnchorBlockHeader);
         });
         break;