refactor: staged writes on note store

yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts

@@ -356,7 +356,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
 
     const pendingNullifiers = this.noteCache.getNullifiers(this.callContext.contractAddress);
 
-    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore);
+    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore, this.jobId);
     const dbNotes = await noteService.getNotes(
       this.callContext.contractAddress,
       owner,

yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts

@@ -257,7 +257,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     offset: number,
     status: NoteStatus,
   ): Promise<NoteData[]> {
-    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore);
+    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore, this.jobId);
 
     const dbNotes = await noteService.getNotes(this.contractAddress, owner, storageSlot, status, this.scopes);
     return pickNotes<NoteData>(dbNotes, {
@@ -359,7 +359,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
 
     await logService.syncTaggedLogs(this.contractAddress, pendingTaggedLogArrayBaseSlot, this.scopes);
 
-    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore);
+    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore, this.jobId);
     await noteService.syncNoteNullifiers(this.contractAddress);
   }
 
@@ -393,7 +393,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       await this.capsuleStore.readCapsuleArray(contractAddress, eventValidationRequestsArrayBaseSlot, this.jobId)
     ).map(EventValidationRequest.fromFields);
 
-    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore);
+    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore, this.jobId);
     const noteDeliveries = noteValidationRequests.map(request =>
       noteService.deliverNote(
         request.contractAddress,

yarn-project/pxe/src/debug/pxe_debug_utils.ts

@@ -1,3 +1,4 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import type { NoteDao, NotesFilter } from '@aztec/stdlib/note';
 
 import type { PXE } from '../pxe.js';
@@ -43,6 +44,6 @@ export class PXEDebugUtils {
     const call = await this.contractStore.getFunctionCall('sync_private_state', [], filter.contractAddress);
     await this.#pxe.simulateUtility(call);
 
-    return this.noteStore.getNotes(filter);
+    return this.noteStore.getNotes(filter, randomBytes(8).toString('hex'));
   }
 }

yarn-project/pxe/src/notes/note_service.test.ts

@@ -45,7 +45,7 @@ describe('NoteService', () => {
     contractAddress = await AztecAddress.random();
 
     // Check that there are no notes in the database
-    const notes = await noteStore.getNotes({ contractAddress });
+    const notes = await noteStore.getNotes({ contractAddress }, 'test');
     expect(notes).toHaveLength(0);
 
     // Check that the expected number of accounts is present
@@ -54,7 +54,7 @@ describe('NoteService', () => {
 
     recipient = await keyStore.addAccount(new Fr(69), Fr.random());
 
-    noteService = new NoteService(noteStore, aztecNode, anchorBlockStore);
+    noteService = new NoteService(noteStore, aztecNode, anchorBlockStore, 'test');
   });
 
   it('should remove notes that have been nullified', async () => {
@@ -66,7 +66,7 @@ describe('NoteService', () => {
     jest.spyOn(noteStore, 'applyNullifiers');
 
     // Add the note to storage
-    await noteStore.addNotes([noteDao], recipient.address);
+    await noteStore.addNotes([noteDao], recipient.address, 'test');
 
     // Set up the nullifier in the merkle tree
     const nullifierIndex = randomDataInBlock(123n);
@@ -76,11 +76,14 @@ describe('NoteService', () => {
     await noteService.syncNoteNullifiers(contractAddress);
 
     // Verify the note was removed by checking storage
-    const remainingNotes = await noteStore.getNotes({
-      contractAddress,
-      status: NoteStatus.ACTIVE,
-      scopes: [recipient.address],
-    });
+    const remainingNotes = await noteStore.getNotes(
+      {
+        contractAddress,
+        status: NoteStatus.ACTIVE,
+        scopes: [recipient.address],
+      },
+      'test',
+    );
     expect(remainingNotes).toHaveLength(0);
 
     // Verify the note was removed by checking the spy
@@ -92,7 +95,7 @@ describe('NoteService', () => {
     const noteDao = await NoteDao.random({ contractAddress });
 
     // Add the note to storage
-    await noteStore.addNotes([noteDao], recipient.address);
+    await noteStore.addNotes([noteDao], recipient.address, 'test');
 
     // No nullifier found in merkle tree
     aztecNode.findLeavesIndexes.mockResolvedValue([undefined]);
@@ -101,11 +104,14 @@ describe('NoteService', () => {
     await noteService.syncNoteNullifiers(contractAddress);
 
     // Verify note still exists
-    const remainingNotes = await noteStore.getNotes({
-      contractAddress,
-      status: NoteStatus.ACTIVE,
-      scopes: [recipient.address],
-    });
+    const remainingNotes = await noteStore.getNotes(
+      {
+        contractAddress,
+        status: NoteStatus.ACTIVE,
+        scopes: [recipient.address],
+      },
+      'test',
+    );
     expect(remainingNotes).toHaveLength(1);
     expect(remainingNotes[0]).toEqual(noteDao);
   });
@@ -118,7 +124,7 @@ describe('NoteService', () => {
     const noteDao = await NoteDao.random({ contractAddress });
 
     // Add the note to storage
-    await noteStore.addNotes([noteDao], recipient.address);
+    await noteStore.addNotes([noteDao], recipient.address, 'test');
 
     // Mock nullifier to only exist after synced block
     aztecNode.findLeavesIndexes.mockImplementation(blockNum => {
@@ -132,11 +138,14 @@ describe('NoteService', () => {
     await noteService.syncNoteNullifiers(contractAddress);
 
     // Verify note still exists
-    const remainingNotes = await noteStore.getNotes({
-      contractAddress,
-      status: NoteStatus.ACTIVE,
-      scopes: [recipient.address],
-    });
+    const remainingNotes = await noteStore.getNotes(
+      {
+        contractAddress,
+        status: NoteStatus.ACTIVE,
+        scopes: [recipient.address],
+      },
+      'test',
+    );
     expect(remainingNotes).toHaveLength(1);
     expect(remainingNotes[0]).toEqual(noteDao);
   });
@@ -157,8 +166,8 @@ describe('NoteService', () => {
     // Verify applyNullifiers was called once for all accounts
     expect(getNotesSpy).toHaveBeenCalledTimes(1);
 
-    // Verify getNotes was called with the correct contract address
-    expect(getNotesSpy).toHaveBeenCalledWith(expect.objectContaining({ contractAddress }));
+    // Verify getNotes was called with the correct contract address and jobId
+    expect(getNotesSpy).toHaveBeenCalledWith(expect.objectContaining({ contractAddress }), 'test');
   });
 
   describe('deliverNote', () => {
@@ -248,7 +257,7 @@ describe('NoteService', () => {
       );
 
       // Verify note was stored
-      const notes = await noteStore.getNotes({ contractAddress, scopes: [recipient.address] });
+      const notes = await noteStore.getNotes({ contractAddress, scopes: [recipient.address] }, 'test');
 
       expect(notes).toHaveLength(1);
       expect(notes[0].noteHash.equals(noteHash)).toBe(true);
@@ -334,19 +343,25 @@ describe('NoteService', () => {
 
       // Now we verify that the note is stored as nullified by checking it can be retrieved only with
       // the ACTIVE_OR_NULLIFIED status on the input.
-      const allNotes = await noteStore.getNotes({
-        contractAddress,
-        scopes: [recipient.address],
-        status: NoteStatus.ACTIVE_OR_NULLIFIED,
-      });
+      const allNotes = await noteStore.getNotes(
+        {
+          contractAddress,
+          scopes: [recipient.address],
+          status: NoteStatus.ACTIVE_OR_NULLIFIED,
+        },
+        'test',
+      );
       expect(allNotes).toHaveLength(1);
       expect(allNotes[0].noteHash.equals(noteHash)).toBe(true);
 
-      const activeNotes = await noteStore.getNotes({
-        contractAddress,
-        scopes: [recipient.address],
-        status: NoteStatus.ACTIVE,
-      });
+      const activeNotes = await noteStore.getNotes(
+        {
+          contractAddress,
+          scopes: [recipient.address],
+          status: NoteStatus.ACTIVE,
+        },
+        'test',
+      );
       expect(activeNotes).toHaveLength(0);
     });
   });

yarn-project/pxe/src/notes/note_service.ts

@@ -15,6 +15,7 @@ export class NoteService {
     private readonly noteStore: NoteStore,
     private readonly aztecNode: AztecNode,
     private readonly anchorBlockStore: AnchorBlockStore,
+    private readonly jobId: string,
   ) {}
 
   /**
@@ -33,13 +34,16 @@ export class NoteService {
     status: NoteStatus,
     scopes?: AztecAddress[],
   ) {
-    const noteDaos = await this.noteStore.getNotes({
-      contractAddress,
-      owner,
-      storageSlot,
-      status,
-      scopes,
-    });
+    const noteDaos = await this.noteStore.getNotes(
+      {
+        contractAddress,
+        owner,
+        storageSlot,
+        status,
+        scopes,
+      },
+      this.jobId,
+    );
     return noteDaos.map(
       ({ contractAddress, owner, storageSlot, randomness, noteNonce, note, noteHash, siloedNullifier }) => ({
         contractAddress,
@@ -70,7 +74,7 @@ export class NoteService {
   public async syncNoteNullifiers(contractAddress: AztecAddress): Promise<void> {
     const syncedBlockNumber = (await this.anchorBlockStore.getBlockHeader()).getBlockNumber();
 
-    const contractNotes = await this.noteStore.getNotes({ contractAddress });
+    const contractNotes = await this.noteStore.getNotes({ contractAddress }, this.jobId);
 
     if (contractNotes.length === 0) {
       return;
@@ -104,7 +108,7 @@ export class NoteService {
       })
       .filter(nullifier => nullifier !== undefined) as DataInBlock<Fr>[];
 
-    await this.noteStore.applyNullifiers(foundNullifiers);
+    await this.noteStore.applyNullifiers(foundNullifiers, this.jobId);
   }
 
   public async deliverNote(
@@ -180,12 +184,17 @@ export class NoteService {
     );
 
     // The note was found by `recipient`, so we use that as the scope when storing the note.
-    await this.noteStore.addNotes([noteDao], recipient);
+    await this.noteStore.addNotes([noteDao], recipient, this.jobId);
 
     if (nullifierIndex !== undefined) {
       // We found nullifier index which implies that the note has already been nullified.
-      const { data: _, ...blockHashAndNum } = nullifierIndex;
-      await this.noteStore.applyNullifiers([{ data: siloedNullifier, ...blockHashAndNum }]);
+      // Only apply nullifier if the note isn't already nullified
+      // (it might have been nullified by a previous sync operation in this job)
+      const alreadyNullified = await this.noteStore.isNoteNullified(noteDao, this.jobId);
+      if (!alreadyNullified) {
+        const { data: _, ...blockHashAndNum } = nullifierIndex;
+        await this.noteStore.applyNullifiers([{ data: siloedNullifier, ...blockHashAndNum }], this.jobId);
+      }
     }
   }
 }

yarn-project/pxe/src/pxe.ts

@@ -158,7 +158,13 @@ export class PXE {
     );
 
     const jobCoordinator = new JobCoordinator(store);
-    jobCoordinator.registerStores([capsuleStore, senderTaggingStore, recipientTaggingStore, privateEventStore]);
+    jobCoordinator.registerStores([
+      capsuleStore,
+      senderTaggingStore,
+      recipientTaggingStore,
+      privateEventStore,
+      noteStore,
+    ]);
 
     const debugUtils = new PXEDebugUtils(contractStore, noteStore);