refactor: crash tolerant PXE

yarn-project/end-to-end/src/e2e_multiple_blobs.test.ts

@@ -58,7 +58,7 @@ describe('e2e_multiple_blobs', () => {
     await aztecNodeAdmin.setConfig({ minTxsPerBlock: TX_COUNT });
 
     const provenTxs = [
-      // 1 contract deployment tx.
+      // 1 contract deployment tx (publishes public_dispatch bytecode: ~1,931 fields).
       await publishContractClass(wallet, AvmTestContract.artifact),
       // 2 private function broadcast txs. We pick [2] because it has large bytecode (~1,807 fields),
       // which combined with the contract class publication exceeds FIELDS_PER_BLOB (4,096).

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

@@ -123,6 +123,7 @@ export class ContractFunctionSimulator {
    * @param senderForTags - The address that is used as a tagging sender when emitting private logs. Returned from
    * the `privateGetSenderForTags` oracle.
    * @param scopes - The accounts whose notes we can access in this call. Currently optional and will default to all.
+   * @param jobId - The job ID for staged writes.
    * @returns The result of the execution.
    */
   public async run(
@@ -131,13 +132,14 @@ export class ContractFunctionSimulator {
     selector: FunctionSelector,
     msgSender = AztecAddress.fromField(Fr.MAX_FIELD_VALUE),
     anchorBlockHeader: BlockHeader,
-    senderForTags?: AztecAddress,
-    scopes?: AztecAddress[],
+    senderForTags: AztecAddress | undefined,
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<PrivateExecutionResult> {
     const simulatorSetupTimer = new Timer();
 
     await this.contractStore.syncPrivateState(contractAddress, selector, privateSyncCall =>
-      this.runUtility(privateSyncCall, [], anchorBlockHeader, scopes),
+      this.runUtility(privateSyncCall, [], anchorBlockHeader, scopes, jobId),
     );
 
     await verifyCurrentClassId(contractAddress, this.aztecNode, this.contractStore, anchorBlockHeader);
@@ -174,7 +176,7 @@ export class ContractFunctionSimulator {
       callContext,
       anchorBlockHeader,
       async call => {
-        await this.runUtility(call, [], anchorBlockHeader, scopes);
+        await this.runUtility(call, [], anchorBlockHeader, scopes, jobId);
       },
       request.authWitnesses,
       request.capsules,
@@ -192,6 +194,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       0, // totalPublicArgsCount
       startSideEffectCounter,
       undefined, // log
@@ -255,13 +258,15 @@ export class ContractFunctionSimulator {
    * @param anchorBlockHeader - The block header to use as base state for this run.
    * @param scopes - Optional array of account addresses whose notes can be accessed in this call. Defaults to all
    * accounts if not specified.
+   * @param jobId - The job ID for staged writes.
    * @returns A return value of the utility function in a form as returned by the simulator (Noir fields)
    */
   public async runUtility(
     call: FunctionCall,
     authwits: AuthWitness[],
     anchorBlockHeader: BlockHeader,
-    scopes?: AztecAddress[],
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<Fr[]> {
     await verifyCurrentClassId(call.to, this.aztecNode, this.contractStore, anchorBlockHeader);
 
@@ -286,6 +291,7 @@ export class ContractFunctionSimulator {
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      jobId,
       undefined,
       scopes,
     );

yarn-project/pxe/src/contract_function_simulator/oracle/oracle_version_is_checked.test.ts

@@ -140,7 +140,16 @@ describe('Oracle Version Check test suite', () => {
       // Call the private function with arbitrary message sender and sender for tags
       const msgSender = await AztecAddress.random();
       const senderForTags = await AztecAddress.random();
-      await acirSimulator.run(txRequest, contractAddress, selector, msgSender, anchorBlockHeader, senderForTags);
+      await acirSimulator.run(
+        txRequest,
+        contractAddress,
+        selector,
+        msgSender,
+        anchorBlockHeader,
+        senderForTags,
+        undefined,
+        'test',
+      );
 
       expect(utilityAssertCompatibleOracleVersionSpy).toHaveBeenCalledTimes(1);
     }, 30_000);
@@ -169,7 +178,7 @@ describe('Oracle Version Check test suite', () => {
       };
 
       // Call the utility function
-      await acirSimulator.runUtility(execRequest, [], anchorBlockHeader, []);
+      await acirSimulator.runUtility(execRequest, [], anchorBlockHeader, [], 'test');
 
       expect(utilityAssertCompatibleOracleVersionSpy).toHaveBeenCalledTimes(1);
     }, 30_000);

yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts

@@ -211,7 +211,16 @@ describe('Private Execution test suite', () => {
       salt: Fr.random(),
     });
 
-    return acirSimulator.run(txRequest, contractAddress, selector, msgSender, anchorBlockHeader, senderForTags);
+    return acirSimulator.run(
+      txRequest,
+      contractAddress,
+      selector,
+      msgSender,
+      anchorBlockHeader,
+      senderForTags,
+      undefined,
+      'test',
+    );
   };
 
   const insertLeaves = async (leaves: Fr[], name = 'noteHash') => {

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

@@ -104,6 +104,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
     senderAddressBookStore: SenderAddressBookStore,
     capsuleStore: CapsuleStore,
     privateEventStore: PrivateEventStore,
+    jobId: string,
     private totalPublicCalldataCount: number = 0,
     protected sideEffectCounter: number = 0,
     log = createLogger('simulator:client_execution_context'),
@@ -126,6 +127,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       senderAddressBookStore,
       capsuleStore,
       privateEventStore,
+      jobId,
       log,
       scopes,
     );
@@ -281,9 +283,9 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       // This is a tagging secret we've not yet used in this tx, so first sync our store to make sure its indices
       // are up to date. We do this here because this store is not synced as part of the global sync because
       // that'd be wasteful as most tagging secrets are not used in each tx.
-      await syncSenderTaggingIndexes(secret, this.contractAddress, this.aztecNode, this.senderTaggingStore);
+      await syncSenderTaggingIndexes(secret, this.contractAddress, this.aztecNode, this.senderTaggingStore, this.jobId);
 
-      const lastUsedIndex = await this.senderTaggingStore.getLastUsedIndex(secret);
+      const lastUsedIndex = await this.senderTaggingStore.getLastUsedIndex(secret, this.jobId);
       // If lastUsedIndex is undefined, we've never used this secret, so start from 0
       // Otherwise, the next index to use is one past the last used index
       return lastUsedIndex === undefined ? 0 : lastUsedIndex + 1;
@@ -372,7 +374,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,
@@ -586,6 +588,7 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP
       this.senderAddressBookStore,
       this.capsuleStore,
       this.privateEventStore,
+      this.jobId,
       this.totalPublicCalldataCount,
       sideEffectCounter,
       this.log,

yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution.test.ts

@@ -190,7 +190,7 @@ describe('Utility Execution test suite', () => {
       returnTypes: artifact.returnTypes,
     };
 
-    const result = await acirSimulator.runUtility(execRequest, [], anchorBlockHeader, []);
+    const result = await acirSimulator.runUtility(execRequest, [], anchorBlockHeader, [], 'test');
 
     expect(result).toEqual([new Fr(9)]);
   }, 30_000);
@@ -222,6 +222,7 @@ describe('Utility Execution test suite', () => {
         senderAddressBookStore,
         capsuleStore,
         privateEventStore,
+        'test',
       );
     });
 

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

@@ -64,6 +64,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     protected readonly senderAddressBookStore: SenderAddressBookStore,
     protected readonly capsuleStore: CapsuleStore,
     protected readonly privateEventStore: PrivateEventStore,
+    protected readonly jobId: string,
     protected log = createLogger('simulator:client_view_context'),
     protected readonly scopes?: AztecAddress[],
   ) {}
@@ -256,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, {
@@ -353,11 +354,12 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       this.recipientTaggingStore,
       this.senderAddressBookStore,
       this.addressStore,
+      this.jobId,
     );
 
     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);
   }
 
@@ -384,14 +386,26 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     // We read all note and event validation requests and process them all concurrently. This makes the process much
     // faster as we don't need to wait for the network round-trip.
     const noteValidationRequests = (
-      await this.capsuleStore.readCapsuleArray(contractAddress, noteValidationRequestsArrayBaseSlot)
+      await this.capsuleStore.readCapsuleArray(contractAddress, noteValidationRequestsArrayBaseSlot, this.jobId)
     ).map(NoteValidationRequest.fromFields);
 
     const eventValidationRequests = (
-      await this.capsuleStore.readCapsuleArray(contractAddress, eventValidationRequestsArrayBaseSlot)
+      await this.capsuleStore.readCapsuleArray(contractAddress, eventValidationRequestsArrayBaseSlot, this.jobId)
     ).map(EventValidationRequest.fromFields);
 
-    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore);
+    this.log.verbose('utilityValidateEnqueuedNotesAndEvents', {
+      contract: contractAddress.toString(),
+      noteValidationRequestCount: noteValidationRequests.length,
+      eventValidationRequestCount: eventValidationRequests.length,
+      notes: noteValidationRequests.map(r => ({
+        owner: r.owner.toString(),
+        recipient: r.recipient.toString(),
+        storageSlot: r.storageSlot.toString(),
+        noteHash: r.noteHash.toString(),
+      })),
+    });
+
+    const noteService = new NoteService(this.noteStore, this.aztecNode, this.anchorBlockStore, this.jobId);
     const noteDeliveries = noteValidationRequests.map(request =>
       noteService.deliverNote(
         request.contractAddress,
@@ -407,7 +421,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       ),
     );
 
-    const eventService = new EventService(this.anchorBlockStore, this.aztecNode, this.privateEventStore);
+    const eventService = new EventService(this.anchorBlockStore, this.aztecNode, this.privateEventStore, this.jobId);
     const eventDeliveries = eventValidationRequests.map(request =>
       eventService.deliverEvent(
         request.contractAddress,
@@ -422,8 +436,8 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     await Promise.all([...noteDeliveries, ...eventDeliveries]);
 
     // Requests are cleared once we're done.
-    await this.capsuleStore.setCapsuleArray(contractAddress, noteValidationRequestsArrayBaseSlot, []);
-    await this.capsuleStore.setCapsuleArray(contractAddress, eventValidationRequestsArrayBaseSlot, []);
+    await this.capsuleStore.setCapsuleArray(contractAddress, noteValidationRequestsArrayBaseSlot, [], this.jobId);
+    await this.capsuleStore.setCapsuleArray(contractAddress, eventValidationRequestsArrayBaseSlot, [], this.jobId);
   }
 
   public async utilityBulkRetrieveLogs(
@@ -439,7 +453,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     // We read all log retrieval requests and process them all concurrently. This makes the process much faster as we
     // don't need to wait for the network round-trip.
     const logRetrievalRequests = (
-      await this.capsuleStore.readCapsuleArray(contractAddress, logRetrievalRequestsArrayBaseSlot)
+      await this.capsuleStore.readCapsuleArray(contractAddress, logRetrievalRequestsArrayBaseSlot, this.jobId)
     ).map(LogRetrievalRequest.fromFields);
 
     const logService = new LogService(
@@ -450,18 +464,20 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       this.recipientTaggingStore,
       this.senderAddressBookStore,
       this.addressStore,
+      this.jobId,
     );
 
     const maybeLogRetrievalResponses = await logService.bulkRetrieveLogs(logRetrievalRequests);
 
     // Requests are cleared once we're done.
-    await this.capsuleStore.setCapsuleArray(contractAddress, logRetrievalRequestsArrayBaseSlot, []);
+    await this.capsuleStore.setCapsuleArray(contractAddress, logRetrievalRequestsArrayBaseSlot, [], this.jobId);
 
     // The responses are stored as Option<LogRetrievalResponse> in a second CapsuleArray.
     await this.capsuleStore.setCapsuleArray(
       contractAddress,
       logRetrievalResponsesArrayBaseSlot,
       maybeLogRetrievalResponses.map(LogRetrievalResponse.toSerializedOption),
+      this.jobId,
     );
   }
 
@@ -470,7 +486,8 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       // TODO(#10727): instead of this check that this.contractAddress is allowed to access the external DB
       throw new Error(`Contract ${contractAddress} is not allowed to access ${this.contractAddress}'s PXE DB`);
     }
-    return this.capsuleStore.storeCapsule(this.contractAddress, slot, capsule);
+    this.capsuleStore.storeCapsule(this.contractAddress, slot, capsule, this.jobId);
+    return Promise.resolve();
   }
 
   public async utilityLoadCapsule(contractAddress: AztecAddress, slot: Fr): Promise<Fr[] | null> {
@@ -481,7 +498,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     return (
       // TODO(#12425): On the following line, the pertinent capsule gets overshadowed by the transient one. Tackle this.
       this.capsules.find(c => c.contractAddress.equals(contractAddress) && c.storageSlot.equals(slot))?.data ??
-      (await this.capsuleStore.loadCapsule(this.contractAddress, slot))
+      (await this.capsuleStore.loadCapsule(this.contractAddress, slot, this.jobId))
     );
   }
 
@@ -490,7 +507,8 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       // TODO(#10727): instead of this check that this.contractAddress is allowed to access the external DB
       throw new Error(`Contract ${contractAddress} is not allowed to access ${this.contractAddress}'s PXE DB`);
     }
-    return this.capsuleStore.deleteCapsule(this.contractAddress, slot);
+    this.capsuleStore.deleteCapsule(this.contractAddress, slot, this.jobId);
+    return Promise.resolve();
   }
 
   public utilityCopyCapsule(
@@ -503,7 +521,7 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
       // TODO(#10727): instead of this check that this.contractAddress is allowed to access the external DB
       throw new Error(`Contract ${contractAddress} is not allowed to access ${this.contractAddress}'s PXE DB`);
     }
-    return this.capsuleStore.copyCapsule(this.contractAddress, srcSlot, dstSlot, numEntries);
+    return this.capsuleStore.copyCapsule(this.contractAddress, srcSlot, dstSlot, numEntries, this.jobId);
   }
 
   // TODO(#11849): consider replacing this oracle with a pure Noir implementation of aes decryption.

yarn-project/pxe/src/events/event_service.test.ts

@@ -30,6 +30,7 @@ describe('deliverEvent', () => {
   let aztecNode: ReturnType<typeof mock<AztecNode>>;
 
   let eventService: EventService;
+  const TEST_JOB_ID = 'test-job-id';
 
   const setSyncedBlockNumber = (blockNumber: BlockNumber) => {
     return anchorBlockStore.setHeader(
@@ -89,7 +90,7 @@ describe('deliverEvent', () => {
       ]),
     );
 
-    eventService = new EventService(anchorBlockStore, aztecNode, privateEventStore);
+    eventService = new EventService(anchorBlockStore, aztecNode, privateEventStore, TEST_JOB_ID);
   });
 
   function runDeliverEvent(
@@ -138,12 +139,16 @@ describe('deliverEvent', () => {
     await runDeliverEvent();
 
     // I should be able to retrieve the private event I just saved using getPrivateEvents
-    const result = await privateEventStore.getPrivateEvents(eventSelector, {
-      contractAddress,
-      fromBlock: blockNumber,
-      toBlock: blockNumber + 1,
-      scopes: [recipient],
-    });
+    const result = await privateEventStore.getPrivateEvents(
+      eventSelector,
+      {
+        contractAddress,
+        fromBlock: blockNumber,
+        toBlock: blockNumber + 1,
+        scopes: [recipient],
+      },
+      TEST_JOB_ID,
+    );
 
     expect(result.length).toEqual(1);
     expect(result[0].packedEvent).toEqual(eventContent);

yarn-project/pxe/src/events/event_service.ts

@@ -14,6 +14,7 @@ export class EventService {
     private readonly anchorBlockStore: AnchorBlockStore,
     private readonly aztecNode: AztecNode,
     private readonly privateEventStore: PrivateEventStore,
+    private readonly jobId: string,
   ) {}
 
   public async deliverEvent(
@@ -72,6 +73,7 @@ export class EventService {
         l2BlockNumber: nullifierIndex.l2BlockNumber, // Block number in which the event was emitted
         l2BlockHash: nullifierIndex.l2BlockHash, // Block hash in which the event was emitted
       },
+      this.jobId,
     );
   }
 }

yarn-project/pxe/src/job_coordinator/index.ts

@@ -0,0 +1 @@
+export { JobCoordinator, type StagedStore } from './job_coordinator.js';