refactor: atomic reorg handling (#19364)

benesjan · web-flow · commit 1596ab942b7d · 2026-01-12T18:03:09.000Z
As Martin mentioned [here](#19327 (comment)) makes sense to have the reorg handling be atomic. In this PR I tackle that. Note that this PR doesn't need to be in stack with #19443 but #19451 depends on both so putting this one into stack was the only way to unblock myself working on #19451.
diff --git a/yarn-project/pxe/src/block_synchronizer/block_synchronizer.test.ts b/yarn-project/pxe/src/block_synchronizer/block_synchronizer.test.ts
@@ -1,5 +1,6 @@
 import { BlockNumber, CheckpointNumber } from '@aztec/foundation/branded-types';
 import { timesParallel } from '@aztec/foundation/collection';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { openTmpStore } from '@aztec/kv-store/lmdb-v2';
 import { L2TipsKVStore } from '@aztec/kv-store/stores';
 import { GENESIS_CHECKPOINT_HEADER_HASH, L2Block, L2BlockNew, type L2BlockStream } from '@aztec/stdlib/block';
@@ -36,7 +37,7 @@ describe('BlockSynchronizer', () => {
     anchorBlockStore = new AnchorBlockStore(store);
     noteStore = await NoteStore.create(store);
     privateEventStore = new PrivateEventStore(store);
-    synchronizer = new TestSynchronizer(aztecNode, anchorBlockStore, noteStore, privateEventStore, tipsStore);
+    synchronizer = new TestSynchronizer(aztecNode, store, anchorBlockStore, noteStore, privateEventStore, tipsStore);
   });
 
   it('sets header from latest block', async () => {
@@ -48,12 +49,14 @@ describe('BlockSynchronizer', () => {
   });
 
   it('removes notes from db on a reorg', async () => {
-    const rollbackNotesAndNullifiers = jest
-      .spyOn(noteStore, 'rollbackNotesAndNullifiers')
-      .mockImplementation(() => Promise.resolve());
-    aztecNode.getBlockHeader.mockImplementation(async blockNumber =>
-      (await L2Block.random(BlockNumber(blockNumber as number))).getBlockHeader(),
-    );
+    const rollback = jest.spyOn(noteStore, 'rollback').mockImplementation(() => Promise.resolve());
+    aztecNode.getBlockHeaderByHash.mockImplementation(async hash => {
+      // For the test, when hash is '0x3', return block header for block 3
+      if (hash.equals(Fr.fromString('0x3'))) {
+        return (await L2Block.random(BlockNumber(3))).getBlockHeader();
+      }
+      return undefined;
+    });
 
     await synchronizer.handleBlockStreamEvent({
       type: 'blocks-added',
@@ -66,16 +69,18 @@ describe('BlockSynchronizer', () => {
       checkpoint: { number: CheckpointNumber.ZERO, hash: GENESIS_CHECKPOINT_HEADER_HASH.toString() },
     });
 
-    expect(rollbackNotesAndNullifiers).toHaveBeenCalledWith(3, 4);
+    expect(rollback).toHaveBeenCalledWith(3, 4);
   });
 
   it('removes private events from db on a reorg', async () => {
-    const rollbackEventsAfterBlock = jest
-      .spyOn(privateEventStore, 'rollbackEventsAfterBlock')
-      .mockImplementation(() => Promise.resolve());
-    aztecNode.getBlockHeader.mockImplementation(async blockNumber =>
-      (await L2Block.random(BlockNumber(blockNumber as number))).getBlockHeader(),
-    );
+    const rollback = jest.spyOn(privateEventStore, 'rollback').mockImplementation(() => Promise.resolve());
+    aztecNode.getBlockHeaderByHash.mockImplementation(async hash => {
+      // For the test, when hash is '0x3', return block header for block 3
+      if (hash.equals(Fr.fromString('0x3'))) {
+        return (await L2Block.random(BlockNumber(3))).getBlockHeader();
+      }
+      return undefined;
+    });
 
     await synchronizer.handleBlockStreamEvent({
       type: 'blocks-added',
@@ -88,6 +93,6 @@ describe('BlockSynchronizer', () => {
       checkpoint: { number: CheckpointNumber.ZERO, hash: GENESIS_CHECKPOINT_HEADER_HASH.toString() },
     });
 
-    expect(rollbackEventsAfterBlock).toHaveBeenCalledWith(3, 4);
+    expect(rollback).toHaveBeenCalledWith(3, 4);
   });
 });
diff --git a/yarn-project/pxe/src/block_synchronizer/block_synchronizer.ts b/yarn-project/pxe/src/block_synchronizer/block_synchronizer.ts
@@ -1,5 +1,7 @@
 import { BlockNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import type { L2TipsKVStore } from '@aztec/kv-store/stores';
 import { L2BlockStream, type L2BlockStreamEvent, type L2BlockStreamEventHandler } from '@aztec/stdlib/block';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
@@ -21,6 +23,7 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
 
   constructor(
     private node: AztecNode,
+    private store: AztecAsyncKVStore,
     private anchorBlockStore: AnchorBlockStore,
     private noteStore: NoteStore,
     private privateEventStore: PrivateEventStore,
@@ -61,17 +64,25 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
       }
       case 'chain-pruned': {
         this.log.warn(`Pruning data after block ${event.block.number} due to reorg`);
-        // We first unnullify and then remove so that unnullified notes that were created after the block number end up deleted.
-        const lastSynchedBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
-        await this.noteStore.rollbackNotesAndNullifiers(event.block.number, lastSynchedBlockNumber);
-        await this.privateEventStore.rollbackEventsAfterBlock(event.block.number, lastSynchedBlockNumber);
-        // Update the header to the last block.
-        const newHeader = await this.node.getBlockHeader(event.block.number);
-        if (!newHeader) {
-          this.log.error(`Block header not found for block number ${event.block.number} during chain prune`);
-        } else {
-          await this.anchorBlockStore.setHeader(newHeader);
+
+        const oldAnchorBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
+        // 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
+        // forward.
+        const newAnchorBlockHeader = await this.node.getBlockHeaderByHash(Fr.fromString(event.block.hash));
+
+        if (!newAnchorBlockHeader) {
+          throw new Error(
+            `Block header for block number ${event.block.number} and hash ${event.block.hash} not found during chain prune. This likely indicates a bug in the node, as we receive block stream events and fetch block headers from the same node.`,
+          );
         }
+
+        // Operations are wrapped in a single transaction to ensure atomicity.
+        await this.store.transactionAsync(async () => {
+          await this.noteStore.rollback(event.block.number, oldAnchorBlockNumber);
+          await this.privateEventStore.rollback(event.block.number, oldAnchorBlockNumber);
+          await this.anchorBlockStore.setHeader(newAnchorBlockHeader);
+        });
         break;
       }
     }
diff --git a/yarn-project/pxe/src/pxe.ts b/yarn-project/pxe/src/pxe.ts
@@ -148,6 +148,7 @@ export class PXE {
     const tipsStore = new L2TipsKVStore(store, 'pxe');
     const synchronizer = new BlockSynchronizer(
       node,
+      store,
       anchorBlockStore,
       noteStore,
       privateEventStore,
diff --git a/yarn-project/pxe/src/storage/note_store/note_store.test.ts b/yarn-project/pxe/src/storage/note_store/note_store.test.ts
@@ -488,7 +488,7 @@ describe('NoteStore', () => {
     });
   });
 
-  describe('NoteStore.rollbackNotesAndNullifiers', () => {
+  describe('NoteStore.rollback', () => {
     let provider: NoteStore;
     let store: AztecLMDBStoreV2;
 
@@ -521,7 +521,7 @@ describe('NoteStore', () => {
         // Apply nullifiers and rollback to block 3
         // - should restore noteBlock3 (nullified at block 4) and preserve noteBlock1 (nullified at block 2)
         await provider.applyNullifiers(nullifiers);
-        await provider.rollbackNotesAndNullifiers(3, 6);
+        await provider.rollback(3, 6);
       }
 
       beforeEach(async () => {
@@ -584,7 +584,7 @@ describe('NoteStore', () => {
 
         // Since nullification happened at block 5 (not after), it should stay nullified
         // The rewind loop processes blocks (blockNumber+1) to synchedBlockNumber = 6 to 5 = no iterations
-        await provider.rollbackNotesAndNullifiers(5, 5);
+        await provider.rollback(5, 5);
 
         const activeNotes = await provider.getNotes({ contractAddress: CONTRACT_A });
         expect(activeNotes).toHaveLength(0);
@@ -610,7 +610,7 @@ describe('NoteStore', () => {
         await provider.applyNullifiers(nullifiers);
 
         // blockNumber=6, synchedBlockNumber=4 therefore no nullifications to rewind
-        await provider.rollbackNotesAndNullifiers(6, 4);
+        await provider.rollback(6, 4);
 
         const activeNotes = await provider.getNotes({ contractAddress: CONTRACT_A });
         expect(activeNotes).toHaveLength(0);
@@ -635,7 +635,7 @@ describe('NoteStore', () => {
           },
         ];
         await provider.applyNullifiers(nullifiers);
-        await provider.rollbackNotesAndNullifiers(5, 100);
+        await provider.rollback(5, 100);
 
         // note1 should be restored (nullified at block 7 > rollback block 5)
         // note2 should be deleted (created at block 10 > rollback block 5)
@@ -644,7 +644,7 @@ describe('NoteStore', () => {
       });
 
       it('handles rollback on empty PXE database gracefully', async () => {
-        await expect(provider.rollbackNotesAndNullifiers(10, 20)).resolves.not.toThrow();
+        await expect(provider.rollback(10, 20)).resolves.not.toThrow();
         const notes = await provider.getNotes({ contractAddress: CONTRACT_A });
         expect(notes).toHaveLength(0);
       });
diff --git a/yarn-project/pxe/src/storage/note_store/note_store.ts b/yarn-project/pxe/src/storage/note_store/note_store.ts
@@ -124,10 +124,12 @@ export class NoteStore {
    * specified block number. It restores any notes that were nullified after the given block
    * and deletes any active notes created after that block.
    *
+   * IMPORTANT: This method must be called within a transaction to ensure atomicity.
+   *
    * @param blockNumber - The new chain tip after a reorg
    * @param synchedBlockNumber - The block number up to which PXE managed to sync before the reorg happened.
    */
-  public async rollbackNotesAndNullifiers(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+  public async rollback(blockNumber: number, synchedBlockNumber: number): Promise<void> {
     await this.#rewindNullifiersAfterBlock(blockNumber, synchedBlockNumber);
     await this.#deleteActiveNotesAfterBlock(blockNumber);
   }
@@ -140,24 +142,22 @@ export class NoteStore {
    *
    * @param blockNumber - Notes created after this block number will be deleted
    */
-  #deleteActiveNotesAfterBlock(blockNumber: number): Promise<void> {
-    return this.#store.transactionAsync(async () => {
-      const notes = await toArray(this.#notes.valuesAsync());
-      for (const note of notes) {
-        const noteDao = NoteDao.fromBuffer(note);
-        if (noteDao.l2BlockNumber > blockNumber) {
-          const noteIndex = toBufferBE(noteDao.index, 32).toString('hex');
-          await this.#notes.delete(noteIndex);
-          await this.#notesToScope.delete(noteIndex);
-          await this.#nullifierToNoteId.delete(noteDao.siloedNullifier.toString());
-          const scopes = await toArray(this.#scopes.keysAsync());
-          for (const scope of scopes) {
-            await this.#notesByContractAndScope.get(scope)!.deleteValue(noteDao.contractAddress.toString(), noteIndex);
-            await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(noteDao.storageSlot.toString(), noteIndex);
-          }
+  async #deleteActiveNotesAfterBlock(blockNumber: number): Promise<void> {
+    const notes = await toArray(this.#notes.valuesAsync());
+    for (const note of notes) {
+      const noteDao = NoteDao.fromBuffer(note);
+      if (noteDao.l2BlockNumber > blockNumber) {
+        const noteIndex = toBufferBE(noteDao.index, 32).toString('hex');
+        await this.#notes.delete(noteIndex);
+        await this.#notesToScope.delete(noteIndex);
+        await this.#nullifierToNoteId.delete(noteDao.siloedNullifier.toString());
+        const scopes = await toArray(this.#scopes.keysAsync());
+        for (const scope of scopes) {
+          await this.#notesByContractAndScope.get(scope)!.deleteValue(noteDao.contractAddress.toString(), noteIndex);
+          await this.#notesByStorageSlotAndScope.get(scope)!.deleteValue(noteDao.storageSlot.toString(), noteIndex);
         }
       }
-    });
+    }
   }
 
   /**
@@ -171,50 +171,52 @@ export class NoteStore {
    * @param synchedBlockNumber - Upper bound for the block range to process
    */
   async #rewindNullifiersAfterBlock(blockNumber: number, synchedBlockNumber: number): Promise<void> {
-    await this.#store.transactionAsync(async () => {
-      const nullifiersToUndo: string[] = [];
-      const currentBlockNumber = blockNumber + 1;
-      for (let i = currentBlockNumber; i <= synchedBlockNumber; i++) {
-        nullifiersToUndo.push(...(await toArray(this.#nullifiersByBlockNumber.getValuesAsync(i))));
-      }
-      const notesIndexesToReinsert = await Promise.all(
-        nullifiersToUndo.map(nullifier => this.#nullifiedNotesByNullifier.getAsync(nullifier)),
-      );
-      const notNullNoteIndexes = notesIndexesToReinsert.filter(noteIndex => noteIndex != undefined);
-      const nullifiedNoteBuffers = await Promise.all(
-        notNullNoteIndexes.map(noteIndex => this.#nullifiedNotes.getAsync(noteIndex!)),
-      );
-      const noteDaos = nullifiedNoteBuffers
-        .filter(buffer => buffer != undefined)
-        .map(buffer => NoteDao.fromBuffer(buffer!));
-
-      for (const dao of noteDaos) {
-        const noteIndex = toBufferBE(dao.index, 32).toString('hex');
-        await this.#notes.set(noteIndex, dao.toBuffer());
-        await this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex);
+    const nullifiersToUndo: string[] = [];
+    const currentBlockNumber = blockNumber + 1;
+    for (let i = currentBlockNumber; i <= synchedBlockNumber; i++) {
+      nullifiersToUndo.push(...(await toArray(this.#nullifiersByBlockNumber.getValuesAsync(i))));
+    }
+    const notesIndexesToReinsert = await Promise.all(
+      nullifiersToUndo.map(nullifier => this.#nullifiedNotesByNullifier.getAsync(nullifier)),
+    );
+    const notNullNoteIndexes = notesIndexesToReinsert.filter(noteIndex => noteIndex != undefined);
+    const nullifiedNoteBuffers = await Promise.all(
+      notNullNoteIndexes.map(noteIndex => this.#nullifiedNotes.getAsync(noteIndex!)),
+    );
+    const noteDaos = nullifiedNoteBuffers
+      .filter(buffer => buffer != undefined)
+      .map(buffer => NoteDao.fromBuffer(buffer!));
 
-        const scopes = await toArray(this.#nullifiedNotesToScope.getValuesAsync(noteIndex));
+    for (const dao of noteDaos) {
+      const noteIndex = toBufferBE(dao.index, 32).toString('hex');
 
-        if (scopes.length === 0) {
-          // We should never run into this error because notes always have a scope assigned to them - either on initial
-          // insertion via `addNotes` or when removing their nullifiers.
-          throw new Error(`No scopes found for nullified note with index ${noteIndex}`);
-        }
+      const scopes = await toArray(this.#nullifiedNotesToScope.getValuesAsync(noteIndex));
 
-        for (const scope of scopes) {
-          await this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex);
-          await this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex);
-          await this.#notesToScope.set(noteIndex, scope);
-        }
+      if (scopes.length === 0) {
+        // We should never run into this error because notes always have a scope assigned to them - either on initial
+        // insertion via `addNotes` or when removing their nullifiers.
+        throw new Error(`No scopes found for nullified note with index ${noteIndex}`);
+      }
 
-        await this.#nullifiedNotes.delete(noteIndex);
-        await this.#nullifiedNotesToScope.delete(noteIndex);
-        await this.#nullifiersByBlockNumber.deleteValue(dao.l2BlockNumber, dao.siloedNullifier.toString());
-        await this.#nullifiedNotesByContract.deleteValue(dao.contractAddress.toString(), noteIndex);
-        await this.#nullifiedNotesByStorageSlot.deleteValue(dao.storageSlot.toString(), noteIndex);
-        await this.#nullifiedNotesByNullifier.delete(dao.siloedNullifier.toString());
+      for (const scope of scopes) {
+        await Promise.all([
+          this.#notesByContractAndScope.get(scope.toString())!.set(dao.contractAddress.toString(), noteIndex),
+          this.#notesByStorageSlotAndScope.get(scope.toString())!.set(dao.storageSlot.toString(), noteIndex),
+          this.#notesToScope.set(noteIndex, scope),
+        ]);
       }
-    });
+
+      await Promise.all([
+        this.#notes.set(noteIndex, dao.toBuffer()),
+        this.#nullifierToNoteId.set(dao.siloedNullifier.toString(), noteIndex),
+        this.#nullifiedNotes.delete(noteIndex),
+        this.#nullifiedNotesToScope.delete(noteIndex),
+        this.#nullifiersByBlockNumber.deleteValue(dao.l2BlockNumber, dao.siloedNullifier.toString()),
+        this.#nullifiedNotesByContract.deleteValue(dao.contractAddress.toString(), noteIndex),
+        this.#nullifiedNotesByStorageSlot.deleteValue(dao.storageSlot.toString(), noteIndex),
+        this.#nullifiedNotesByNullifier.delete(dao.siloedNullifier.toString()),
+      ]);
+    }
   }
 
   /**
diff --git a/yarn-project/pxe/src/storage/private_event_store/private_event_store.test.ts b/yarn-project/pxe/src/storage/private_event_store/private_event_store.test.ts
@@ -286,7 +286,7 @@ describe('PrivateEventStore', () => {
       });
 
       // Rollback to block 150 (should remove events from blocks 200 and 300)
-      await privateEventStore.rollbackEventsAfterBlock(150, 300);
+      await privateEventStore.rollback(150, 300);
 
       const events = await privateEventStore.getPrivateEvents(eventSelector, {
         contractAddress,
@@ -315,7 +315,7 @@ describe('PrivateEventStore', () => {
       });
 
       // Rollback to block 100
-      await privateEventStore.rollbackEventsAfterBlock(100, 200);
+      await privateEventStore.rollback(100, 200);
 
       // Verify event was removed
       let events = await privateEventStore.getPrivateEvents(eventSelector, {
@@ -355,7 +355,7 @@ describe('PrivateEventStore', () => {
       });
 
       // Rollback after all existing events
-      await privateEventStore.rollbackEventsAfterBlock(200, 300);
+      await privateEventStore.rollback(200, 300);
 
       const events = await privateEventStore.getPrivateEvents(eventSelector, {
         contractAddress,
diff --git a/yarn-project/pxe/src/storage/private_event_store/private_event_store.ts b/yarn-project/pxe/src/storage/private_event_store/private_event_store.ts
diff --git a/yarn-project/stdlib/src/block/l2_block_stream/interfaces.ts b/yarn-project/stdlib/src/block/l2_block_stream/interfaces.ts
