@aztec/pxe 0.0.1-commit.6c91f13 → 0.0.1-commit.6d3c34e

package/src/pxe.ts

@@ -60,9 +60,11 @@ import {
 } from './contract_function_simulator/contract_function_simulator.js';
 import { readCurrentClassId } from './contract_function_simulator/oracle/private_execution.js';
 import { ProxiedContractStoreFactory } from './contract_function_simulator/proxied_contract_data_source.js';
+import { ProxiedNodeFactory } from './contract_function_simulator/proxied_node.js';
 import { PXEDebugUtils } from './debug/pxe_debug_utils.js';
 import { enrichPublicSimulationError, enrichSimulationError } from './error_enriching.js';
 import { PrivateEventFilterValidator } from './events/private_event_filter_validator.js';
+import { JobCoordinator } from './job_coordinator/job_coordinator.js';
 import {
   PrivateKernelExecutionProver,
   type PrivateKernelExecutionProverConfig,
@@ -107,6 +109,7 @@ export class PXE {
     private protocolContractsProvider: ProtocolContractsProvider,
     private log: Logger,
     private jobQueue: SerialQueue,
+    private jobCoordinator: JobCoordinator,
     public debug: PXEDebugUtils,
   ) {}
 
@@ -143,7 +146,17 @@ export class PXE {
     const capsuleStore = new CapsuleStore(store);
     const keyStore = new KeyStore(store);
     const tipsStore = new L2TipsKVStore(store, 'pxe');
-    const synchronizer = new BlockSynchronizer(node, anchorBlockStore, noteStore, tipsStore, config, loggerOrSuffix);
+    const synchronizer = new BlockSynchronizer(
+      node,
+      anchorBlockStore,
+      noteStore,
+      privateEventStore,
+      tipsStore,
+      config,
+      loggerOrSuffix,
+    );
+
+    const jobCoordinator = new JobCoordinator(store);
 
     const debugUtils = new PXEDebugUtils(contractStore, noteStore);
 
@@ -168,6 +181,7 @@ export class PXE {
       protocolContractsProvider,
       log,
       jobQueue,
+      jobCoordinator,
       debugUtils,
     );
 
@@ -191,7 +205,7 @@ export class PXE {
       this.noteStore,
       this.keyStore,
       this.addressStore,
-      this.node,
+      ProxiedNodeFactory.create(this.node),
       this.anchorBlockStore,
       this.senderTaggingStore,
       this.recipientTaggingStore,
@@ -222,7 +236,7 @@ export class PXE {
    *
    * Useful for tasks that cannot run concurrently, such as contract function simulation.
    */
-  #putInJobQueue<T>(fn: () => Promise<T>): Promise<T> {
+  #putInJobQueue<T>(fn: (jobId: string) => Promise<T>): Promise<T> {
     // TODO(#12636): relax the conditions under which we forbid concurrency.
     if (this.jobQueue.length() != 0) {
       this.log.warn(
@@ -230,7 +244,22 @@ export class PXE {
       );
     }
 
-    return this.jobQueue.put(fn);
+    return this.jobQueue.put(async () => {
+      const jobId = this.jobCoordinator.beginJob();
+      this.log.verbose(`Beginning job ${jobId}`);
+
+      try {
+        const result = await fn(jobId);
+        this.log.verbose(`Committing job ${jobId}`);
+
+        await this.jobCoordinator.commitJob(jobId);
+        return result;
+      } catch (err) {
+        this.log.verbose(`Aborting job ${jobId}`);
+        await this.jobCoordinator.abortJob(jobId);
+        throw err;
+      }
+    });
   }
 
   async #registerProtocolContracts() {
@@ -263,7 +292,8 @@ export class PXE {
   async #executePrivate(
     contractFunctionSimulator: ContractFunctionSimulator,
     txRequest: TxExecutionRequest,
-    scopes?: AztecAddress[],
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ): Promise<PrivateExecutionResult> {
     const { origin: contractAddress, functionSelector } = txRequest;
 
@@ -280,6 +310,7 @@ export class PXE {
         // contract entrypoint
         undefined, // senderForTags
         scopes,
+        jobId,
       );
       this.log.debug(`Private simulation completed for ${contractAddress.toString()}:${functionSelector}`);
       return result;
@@ -298,17 +329,19 @@ export class PXE {
    * @param authWitnesses - Authentication witnesses required for the function call.
    * @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 The simulation result containing the outputs of the utility function.
    */
   async #simulateUtility(
     contractFunctionSimulator: ContractFunctionSimulator,
     call: FunctionCall,
-    authWitnesses?: AuthWitness[],
-    scopes?: AztecAddress[],
+    authWitnesses: AuthWitness[] | undefined,
+    scopes: AztecAddress[] | undefined,
+    jobId: string,
   ) {
     try {
       const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader();
-      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes);
+      return contractFunctionSimulator.runUtility(call, authWitnesses ?? [], anchorBlockHeader, scopes, jobId);
     } catch (err) {
       if (err instanceof SimulationError) {
         await enrichSimulationError(err, this.contractStore, this.log);
@@ -657,14 +690,14 @@ export class PXE {
     let privateExecutionResult: PrivateExecutionResult;
     // We disable proving concurrently mostly out of caution, since it accesses some of our stores. Proving is so
     // computationally demanding that it'd be rare for someone to try to do it concurrently regardless.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const syncTimer = new Timer();
         await this.blockStateSynchronizer.sync();
         const syncTime = syncTimer.ms();
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, undefined, jobId);
 
         const {
           publicInputs,
@@ -741,7 +774,7 @@ export class PXE {
     skipProofGeneration: boolean = true,
   ): Promise<TxProfileResult> {
     // We disable concurrent profiles for consistency with simulateTx.
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
       try {
         const txInfo = {
@@ -761,7 +794,12 @@ export class PXE {
         const syncTime = syncTimer.ms();
 
         const contractFunctionSimulator = this.#getSimulatorForTx();
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest);
+        const privateExecutionResult = await this.#executePrivate(
+          contractFunctionSimulator,
+          txRequest,
+          undefined,
+          jobId,
+        );
 
         const { executionSteps, timings: { proving } = {} } = await this.#prove(
           txRequest,
@@ -841,7 +879,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const txInfo = {
@@ -867,7 +905,7 @@ export class PXE {
         const skipKernels = overrides?.contracts !== undefined && Object.keys(overrides.contracts ?? {}).length > 0;
 
         // Execution of private functions only; no proving, and no kernel logic.
-        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes);
+        const privateExecutionResult = await this.#executePrivate(contractFunctionSimulator, txRequest, scopes, jobId);
 
         let publicInputs: PrivateKernelTailCircuitPublicInputs | undefined;
         let executionSteps: PrivateExecutionStep[] = [];
@@ -982,7 +1020,7 @@ export class PXE {
     // We disable concurrent simulations since those might execute oracles which read and write to the PXE stores (e.g.
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       try {
         const totalTimer = new Timer();
         const syncTimer = new Timer();
@@ -992,10 +1030,16 @@ export class PXE {
         const contractFunctionSimulator = this.#getSimulatorForTx();
 
         await this.contractStore.syncPrivateState(call.to, call.selector, privateSyncCall =>
-          this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+          this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
         );
 
-        const executionResult = await this.#simulateUtility(contractFunctionSimulator, call, authwits ?? [], scopes);
+        const executionResult = await this.#simulateUtility(
+          contractFunctionSimulator,
+          call,
+          authwits ?? [],
+          scopes,
+          jobId,
+        );
         const functionTime = functionTimer.ms();
 
         const totalTime = totalTimer.ms();
@@ -1037,14 +1081,15 @@ export class PXE {
    * @returns - The packed events with block and tx metadata.
    */
   public getPrivateEvents(eventSelector: EventSelector, filter: PrivateEventFilter): Promise<PackedPrivateEvent[]> {
-    return this.#putInJobQueue(async () => {
+    return this.#putInJobQueue(async jobId => {
       await this.blockStateSynchronizer.sync();
       const contractFunctionSimulator = this.#getSimulatorForTx();
 
       await this.contractStore.syncPrivateState(
         filter.contractAddress,
         null,
-        async privateSyncCall => await this.#simulateUtility(contractFunctionSimulator, privateSyncCall),
+        async privateSyncCall =>
+          await this.#simulateUtility(contractFunctionSimulator, privateSyncCall, [], undefined, jobId),
       );
 
       const sanitizedFilter = await new PrivateEventFilterValidator(this.anchorBlockStore).validate(filter);

package/src/storage/private_event_store/private_event_store.ts

@@ -2,7 +2,7 @@ import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { createLogger } from '@aztec/foundation/log';
 import { BufferReader, serializeToBuffer } from '@aztec/foundation/serialize';
-import type { AztecAsyncArray, AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
+import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { EventSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { L2BlockHash } from '@aztec/stdlib/block';
@@ -19,11 +19,14 @@ export type PrivateEventStoreFilter = {
 };
 
 type PrivateEventEntry = {
+  randomness: Fr; // Note that this value is currently not being returned on queries and is therefore temporarily unused
   msgContent: Buffer;
   eventCommitmentIndex: number;
   l2BlockNumber: number;
   l2BlockHash: Buffer;
   txHash: Buffer;
+  /** The lookup key for #eventsByContractScopeSelector, used for cleanup during rollback */
+  lookupKey: string;
 };
 
 type PrivateEventMetadata = InTx & {
@@ -36,10 +39,12 @@ type PrivateEventMetadata = InTx & {
  */
 export class PrivateEventStore {
   #store: AztecAsyncKVStore;
-  /** Array storing the actual private event log entries containing the log content and block number */
-  #eventLogs: AztecAsyncArray<PrivateEventEntry>;
-  /** Map from contract_address_scope_eventSelector to array of indices into #eventLogs for efficient lookup */
-  #eventLogIndex: AztecAsyncMap<string, number[]>;
+  /** Map storing the actual private event log entries, keyed by eventCommitmentIndex */
+  #eventLogs: AztecAsyncMap<number, PrivateEventEntry>;
+  /** Map from contractAddress_scope_eventSelector to eventCommitmentIndex[] for efficient lookup */
+  #eventsByContractScopeSelector: AztecAsyncMap<string, number[]>;
+  /** Map from block number to eventCommitmentIndex[] for rollback support */
+  #eventsByBlockNumber: AztecAsyncMap<number, number[]>;
   /** Map from eventCommitmentIndex to boolean indicating if log has been seen. */
   #seenLogs: AztecAsyncMap<number, boolean>;
 
@@ -47,9 +52,10 @@ export class PrivateEventStore {
 
   constructor(store: AztecAsyncKVStore) {
     this.#store = store;
-    this.#eventLogs = this.#store.openArray('private_event_logs');
-    this.#eventLogIndex = this.#store.openMap('private_event_log_index');
+    this.#eventLogs = this.#store.openMap('private_event_logs');
+    this.#eventsByContractScopeSelector = this.#store.openMap('events_by_contract_scope_selector');
     this.#seenLogs = this.#store.openMap('seen_logs');
+    this.#eventsByBlockNumber = this.#store.openMap('events_by_block_number');
   }
 
   #keyFor(contractAddress: AztecAddress, scope: AztecAddress, eventSelector: EventSelector): string {
@@ -69,6 +75,7 @@ export class PrivateEventStore {
    */
   storePrivateEventLog(
     eventSelector: EventSelector,
+    randomness: Fr,
     msgContent: Fr[],
     eventCommitmentIndex: number,
     metadata: PrivateEventMetadata,
@@ -87,17 +94,21 @@ export class PrivateEventStore {
 
       this.logger.verbose('storing private event log', { contractAddress, scope, msgContent, l2BlockNumber });
 
-      const index = await this.#eventLogs.lengthAsync();
-      await this.#eventLogs.push({
+      await this.#eventLogs.set(eventCommitmentIndex, {
+        randomness,
         msgContent: serializeToBuffer(msgContent),
         l2BlockNumber,
         l2BlockHash: l2BlockHash.toBuffer(),
         eventCommitmentIndex,
         txHash: txHash.toBuffer(),
+        lookupKey: key,
       });
 
-      const existingIndices = (await this.#eventLogIndex.getAsync(key)) || [];
-      await this.#eventLogIndex.set(key, [...existingIndices, index]);
+      const existingIndices = (await this.#eventsByContractScopeSelector.getAsync(key)) || [];
+      await this.#eventsByContractScopeSelector.set(key, [...existingIndices, eventCommitmentIndex]);
+
+      const existingBlockIndices = (await this.#eventsByBlockNumber.getAsync(l2BlockNumber)) || [];
+      await this.#eventsByBlockNumber.set(l2BlockNumber, [...existingBlockIndices, eventCommitmentIndex]);
 
       // Mark this log as seen using eventCommitmentIndex
       await this.#seenLogs.set(eventCommitmentIndex, true);
@@ -123,10 +134,10 @@ export class PrivateEventStore {
 
     for (const scope of filter.scopes) {
       const key = this.#keyFor(filter.contractAddress, scope, eventSelector);
-      const indices = (await this.#eventLogIndex.getAsync(key)) || [];
+      const eventCommitmentIndices = (await this.#eventsByContractScopeSelector.getAsync(key)) || [];
 
-      for (const index of indices) {
-        const entry = await this.#eventLogs.atAsync(index);
+      for (const eventCommitmentIndex of eventCommitmentIndices) {
+        const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
         if (!entry || entry.l2BlockNumber < filter.fromBlock || entry.l2BlockNumber >= filter.toBlock) {
           continue;
         }
@@ -159,4 +170,47 @@ export class PrivateEventStore {
     events.sort((a, b) => a.eventCommitmentIndex - b.eventCommitmentIndex);
     return events.map(ev => ev.event);
   }
+
+  /**
+   * Rolls back private events that were stored after a given `blockNumber` and up to `synchedBlockNumber` (the block
+   * number up to which PXE managed to sync before the reorg happened).
+   */
+  public async rollbackEventsAfterBlock(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+    await this.#store.transactionAsync(async () => {
+      let removedCount = 0;
+
+      for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
+        const indices = await this.#eventsByBlockNumber.getAsync(block);
+        if (indices) {
+          await this.#eventsByBlockNumber.delete(block);
+
+          for (const eventCommitmentIndex of indices) {
+            const entry = await this.#eventLogs.getAsync(eventCommitmentIndex);
+            if (!entry) {
+              throw new Error(`Event log not found for eventCommitmentIndex ${eventCommitmentIndex}`);
+            }
+
+            await this.#eventLogs.delete(eventCommitmentIndex);
+            await this.#seenLogs.delete(eventCommitmentIndex);
+
+            // Update #eventsByContractScopeSelector using the stored lookupKey
+            const existingIndices = await this.#eventsByContractScopeSelector.getAsync(entry.lookupKey);
+            if (!existingIndices || existingIndices.length === 0) {
+              throw new Error(`No indices found in #eventsByContractScopeSelector for key ${entry.lookupKey}`);
+            }
+            const filteredIndices = existingIndices.filter(idx => idx !== eventCommitmentIndex);
+            if (filteredIndices.length === 0) {
+              await this.#eventsByContractScopeSelector.delete(entry.lookupKey);
+            } else {
+              await this.#eventsByContractScopeSelector.set(entry.lookupKey, filteredIndices);
+            }
+
+            removedCount++;
+          }
+        }
+      }
+
+      this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
+    });
+  }
 }

package/src/storage/tagging_store/sender_tagging_store.ts

@@ -3,7 +3,7 @@ import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { DirectionalAppTaggingSecret, PreTag } from '@aztec/stdlib/logs';
 import { TxHash } from '@aztec/stdlib/tx';
 
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../../tagging/index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../../tagging/constants.js';
 
 /**
  * Data provider of tagging data used when syncing the sender tagging indexes. The recipient counterpart of this class

package/src/tagging/constants.ts

@@ -0,0 +1,10 @@
+// This window has to be as large as the largest expected number of logs emitted in a tx for a given directional app
+// tagging secret. If we get more tag indexes consumed than this window, an error is thrown in `PXE::proveTx` function.
+// This is set to a larger value than MAX_PRIVATE_LOGS_PER_TX (currently 64) because there could be more than
+// MAX_PRIVATE_LOGS_PER_TX indexes consumed in case the logs are squashed. This happens when the log contains a note
+// and the note is nullified in the same tx.
+//
+// Having a large window significantly slowed down `e2e_l1_with_wall_time` test as there we perform sync for more than
+// 1000 secrets. For this reason we set it to a relatively low value of 20. 20 should be sufficient for all the use
+// cases.
+export const UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN = 20;

package/src/tagging/index.ts

@@ -11,17 +11,7 @@
 
 export { loadPrivateLogsForSenderRecipientPair } from './recipient_sync/load_private_logs_for_sender_recipient_pair.js';
 export { syncSenderTaggingIndexes } from './sender_sync/sync_sender_tagging_indexes.js';
-
-// This window has to be as large as the largest expected number of logs emitted in a tx for a given directional app
-// tagging secret. If we get more tag indexes consumed than this window, an error is thrown in `PXE::proveTx` function.
-// This is set to a larger value than MAX_PRIVATE_LOGS_PER_TX (currently 64) because there could be more than
-// MAX_PRIVATE_LOGS_PER_TX indexes consumed in case the logs are squashed. This happens when the log contains a note
-// and the note is nullified in the same tx.
-//
-// Having a large window significantly slowed down `e2e_l1_with_wall_time` test as there we perform sync for more than
-// 1000 secrets. For this reason we set it to a relatively low value of 20. 20 should be sufficient for all the use
-// cases.
-export const UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN = 20;
+export { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from './constants.js';
 
 // Re-export tagging-related types from stdlib
 export { DirectionalAppTaggingSecret, Tag, SiloedTag } from '@aztec/stdlib/logs';

package/src/tagging/recipient_sync/load_private_logs_for_sender_recipient_pair.ts

@@ -4,7 +4,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { DirectionalAppTaggingSecret, TxScopedL2Log } from '@aztec/stdlib/logs';
 
 import type { RecipientTaggingStore } from '../../storage/tagging_store/recipient_tagging_store.js';
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../constants.js';
 import { findHighestIndexes } from './utils/find_highest_indexes.js';
 import { loadLogsForRange } from './utils/load_logs_for_range.js';
 
@@ -68,7 +68,10 @@ export async function loadPrivateLogsForSenderRecipientPair(
       throw new Error('Node failed to return latest block header when syncing logs');
     }
 
-    [finalizedBlockNumber, currentTimestamp] = [l2Tips.finalized.number, latestBlockHeader.globalVariables.timestamp];
+    [finalizedBlockNumber, currentTimestamp] = [
+      l2Tips.finalized.block.number,
+      latestBlockHeader.globalVariables.timestamp,
+    ];
   }
 
   let start: number, end: number;

package/src/tagging/sender_sync/sync_sender_tagging_indexes.ts

@@ -3,7 +3,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/server';
 import type { DirectionalAppTaggingSecret } from '@aztec/stdlib/logs';
 
 import type { SenderTaggingStore } from '../../storage/tagging_store/sender_tagging_store.js';
-import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../index.js';
+import { UNFINALIZED_TAGGING_INDEXES_WINDOW_LEN } from '../constants.js';
 import { getStatusChangeOfPending } from './utils/get_status_change_of_pending.js';
 import { loadAndStoreNewTaggingIndexes } from './utils/load_and_store_new_tagging_indexes.js';
 

package/src/tagging/sender_sync/utils/get_status_change_of_pending.ts

@@ -10,7 +10,7 @@ export async function getStatusChangeOfPending(
   aztecNode: AztecNode,
 ): Promise<{ txHashesToFinalize: TxHash[]; txHashesToDrop: TxHash[] }> {
   // Get receipts for all pending tx hashes and the finalized block number.
-  const [receipts, { finalized }] = await Promise.all([
+  const [receipts, tips] = await Promise.all([
     Promise.all(pending.map(pendingTxHash => aztecNode.getTxReceipt(pendingTxHash))),
     aztecNode.getL2Tips(),
   ]);
@@ -22,7 +22,11 @@ export async function getStatusChangeOfPending(
     const receipt = receipts[i];
     const txHash = pending[i];
 
-    if (receipt.status === TxStatus.SUCCESS && receipt.blockNumber && receipt.blockNumber <= finalized.number) {
+    if (
+      receipt.status === TxStatus.SUCCESS &&
+      receipt.blockNumber &&
+      receipt.blockNumber <= tips.finalized.block.number
+    ) {
       // Tx has been included in a block and the corresponding block is finalized --> we mark the indexes as
       // finalized.
       txHashesToFinalize.push(txHash);