@aztec/pxe 0.0.1-commit.c7c42ec → 0.0.1-commit.f295ac2

package/src/storage/private_event_store/private_event_store.ts

@@ -2,12 +2,13 @@ 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, AztecAsyncMultiMap } 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';
 import { type InTx, TxHash } from '@aztec/stdlib/tx';
 
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
 import type { PackedPrivateEvent } from '../../pxe.js';
 
 export type PrivateEventStoreFilter = {
@@ -19,48 +20,141 @@ 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 index of the tx within the block, used for ordering events */
+  txIndexInBlock: number;
+  /** The index of the event within the tx (based on nullifier position), used for ordering events */
+  eventIndexInTx: number;
+  /** The lookup key for #eventsByContractScopeSelector, used for cleanup during rollback */
+  lookupKey: string;
 };
 
 type PrivateEventMetadata = InTx & {
   contractAddress: AztecAddress;
   scope: AztecAddress;
+  /** The index of the tx within the block */
+  txIndexInBlock: number;
+  /** The index of the event within the tx (based on nullifier position) */
+  eventIndexInTx: number;
 };
 
 /**
  * Stores decrypted private event logs.
  */
-export class PrivateEventStore {
+export class PrivateEventStore implements StagedStore {
+  readonly storeName: string = 'private_event';
+
   #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 from eventCommitmentIndex to boolean indicating if log has been seen. */
-  #seenLogs: AztecAsyncMap<number, boolean>;
+  /** Map storing the actual private event log entries, keyed by siloedEventCommitment */
+  #eventLogs: AztecAsyncMap<string, PrivateEventEntry>;
+  /** Multi-map from contractAddress_scope_eventSelector to siloedEventCommitment for efficient lookup */
+  #eventsByContractScopeSelector: AztecAsyncMultiMap<string, string>;
+  /** Multi-map from block number to siloedEventCommitment for rollback support */
+  #eventsByBlockNumber: AztecAsyncMultiMap<number, string>;
+  /** Map from siloedEventCommitment to boolean indicating if log has been seen. */
+  #seenLogs: AztecAsyncMap<string, boolean>;
+
+  /** jobId => eventId (event siloed nullifier)  => PrivateEventEntry */
+  #eventLogsInJobStage: Map<string, Map<string, PrivateEventEntry>>;
 
   logger = createLogger('private_event_store');
 
   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.openMultiMap('events_by_contract_scope_selector');
     this.#seenLogs = this.#store.openMap('seen_logs');
+    this.#eventsByBlockNumber = this.#store.openMultiMap('events_by_block_number');
+
+    this.#eventLogsInJobStage = new Map();
   }
 
   #keyFor(contractAddress: AztecAddress, scope: AztecAddress, eventSelector: EventSelector): string {
     return `${contractAddress.toString()}_${scope.toString()}_${eventSelector.toString()}`;
   }
 
+  async commit(jobId: string): Promise<void> {
+    await Promise.all(
+      [...this.#getEventLogsInJobStage(jobId).entries()].map(async ([eventId, eventEntry]) => {
+        this.logger.verbose('storing private event log (KV store)', eventEntry);
+
+        await Promise.all([
+          this.#eventLogs.set(eventId, eventEntry),
+          this.#eventsByContractScopeSelector.set(eventEntry.lookupKey, eventId),
+          this.#eventsByBlockNumber.set(eventEntry.l2BlockNumber, eventId),
+          this.#seenLogs.set(eventId, true),
+        ]);
+      }),
+    );
+
+    return this.discardStaged(jobId);
+  }
+
+  discardStaged(jobId: string): Promise<void> {
+    this.#eventLogsInJobStage.delete(jobId);
+    return Promise.resolve();
+  }
+
+  #getEventLogsInJobStage(jobId: string): Map<string, PrivateEventEntry> {
+    let jobStage = this.#eventLogsInJobStage.get(jobId);
+    if (jobStage === undefined) {
+      jobStage = new Map();
+      this.#eventLogsInJobStage.set(jobId, jobStage);
+    }
+    return jobStage;
+  }
+
+  async #isSeenLog(jobId: string, eventId: string): Promise<boolean> {
+    const eventLogsInJobStage = this.#getEventLogsInJobStage(jobId).get(eventId);
+    return !!eventLogsInJobStage || !!(await this.#seenLogs.getAsync(eventId));
+  }
+
+  #addEventLogToStage(jobId: string, eventId: string, eventEntry: PrivateEventEntry) {
+    this.#getEventLogsInJobStage(jobId).set(eventId, eventEntry);
+  }
+
+  async #getEventSiloedNullifiers(
+    contractAddress: AztecAddress,
+    scope: AztecAddress,
+    eventSelector: EventSelector,
+    jobId?: string,
+  ): Promise<string[]> {
+    const key = this.#keyFor(contractAddress, scope, eventSelector);
+    const eventSiloedNullifiersInStorage: string[] = [];
+    for await (const eventId of this.#eventsByContractScopeSelector.getValuesAsync(key)) {
+      eventSiloedNullifiersInStorage.push(eventId);
+    }
+
+    if (!jobId) {
+      return eventSiloedNullifiersInStorage;
+    }
+
+    const eventSiloedNullifiersInJobStage = new Set(
+      [...this.#getEventLogsInJobStage(jobId).entries()]
+        .filter(([_, entry]) => entry.lookupKey === key)
+        .map(([idx, _]) => idx),
+    );
+    return [...new Set(eventSiloedNullifiersInStorage).union(eventSiloedNullifiersInJobStage)];
+  }
+
+  async #getEventLogBySiloedNullifier(eventId: string, jobId?: string): Promise<PrivateEventEntry | undefined> {
+    if (jobId) {
+      return this.#getEventLogsInJobStage(jobId).get(eventId) ?? (await this.#eventLogs.getAsync(eventId));
+    } else {
+      return await this.#eventLogs.getAsync(eventId);
+    }
+  }
+
   /**
    * Store a private event log.
    * @param eventSelector - The event selector of the event.
+   * @param randomness - The randomness used for the event commitment.
    * @param msgContent - The content of the event.
-   * @param eventCommitmentIndex - The index of the event commitment in the nullifier tree.
+   * @param siloedEventCommitment - The siloed event commitment (used as unique identifier).
    * @param metadata
    *  contractAddress - The address of the contract that emitted the event.
    *  scope - The address to which the event is scoped.
@@ -69,38 +163,44 @@ export class PrivateEventStore {
    */
   storePrivateEventLog(
     eventSelector: EventSelector,
+    randomness: Fr,
     msgContent: Fr[],
-    eventCommitmentIndex: number,
+    siloedEventCommitment: Fr,
     metadata: PrivateEventMetadata,
+    jobId: string,
   ): Promise<void> {
-    const { contractAddress, scope, txHash, l2BlockNumber, l2BlockHash } = metadata;
+    const { contractAddress, scope, txHash, l2BlockNumber, l2BlockHash, txIndexInBlock, eventIndexInTx } = metadata;
 
     return this.#store.transactionAsync(async () => {
       const key = this.#keyFor(contractAddress, scope, eventSelector);
 
-      // Check if this exact log has already been stored using eventCommitmentIndex as unique identifier
-      const hasBeenSeen = await this.#seenLogs.getAsync(eventCommitmentIndex);
+      // The siloed event commitment is guaranteed to be unique as it's inserted into the nullifier tree. For this
+      // reason we use it as id.
+      const eventId = siloedEventCommitment.toString();
+
+      const hasBeenSeen = await this.#isSeenLog(jobId, eventId);
       if (hasBeenSeen) {
-        this.logger.verbose('Ignoring duplicate event log', { txHash: txHash.toString(), eventCommitmentIndex });
+        this.logger.verbose('Ignoring duplicate event log', { txHash: txHash.toString(), siloedEventCommitment });
         return;
       }
 
-      this.logger.verbose('storing private event log', { contractAddress, scope, msgContent, l2BlockNumber });
+      this.logger.verbose('storing private event log (job stage)', {
+        contractAddress,
+        scope,
+        msgContent,
+        l2BlockNumber,
+      });
 
-      const index = await this.#eventLogs.lengthAsync();
-      await this.#eventLogs.push({
+      this.#addEventLogToStage(jobId, eventId, {
+        randomness,
         msgContent: serializeToBuffer(msgContent),
         l2BlockNumber,
         l2BlockHash: l2BlockHash.toBuffer(),
-        eventCommitmentIndex,
         txHash: txHash.toBuffer(),
+        txIndexInBlock,
+        eventIndexInTx,
+        lookupKey: key,
       });
-
-      const existingIndices = (await this.#eventLogIndex.getAsync(key)) || [];
-      await this.#eventLogIndex.set(key, [...existingIndices, index]);
-
-      // Mark this log as seen using eventCommitmentIndex
-      await this.#seenLogs.set(eventCommitmentIndex, true);
     });
   }
 
@@ -112,21 +212,26 @@ export class PrivateEventStore {
    *  fromBlock: The block number to search from (inclusive).
    *  toBlock: The block number to search upto (exclusive).
    *  scope: - The addresses that decrypted the logs.
-   * @returns - The event log contents, augmented with metadata about
-   *  the transaction and block it the event was included in .
+   * @returns - The event log contents, augmented with metadata about the transaction and block in which the event was
+   * included.
    */
   public async getPrivateEvents(
     eventSelector: EventSelector,
     filter: PrivateEventStoreFilter,
   ): Promise<PackedPrivateEvent[]> {
-    const events: Array<{ eventCommitmentIndex: number; event: PackedPrivateEvent }> = [];
+    const events: Array<{
+      l2BlockNumber: number;
+      txIndexInBlock: number;
+      eventIndexInTx: number;
+      event: PackedPrivateEvent;
+    }> = [];
 
     for (const scope of filter.scopes) {
-      const key = this.#keyFor(filter.contractAddress, scope, eventSelector);
-      const indices = (await this.#eventLogIndex.getAsync(key)) || [];
+      const eventIds = await this.#getEventSiloedNullifiers(filter.contractAddress, scope, eventSelector);
+
+      for (const eventId of eventIds) {
+        const entry = await this.#getEventLogBySiloedNullifier(eventId);
 
-      for (const index of indices) {
-        const entry = await this.#eventLogs.atAsync(index);
         if (!entry || entry.l2BlockNumber < filter.fromBlock || entry.l2BlockNumber >= filter.toBlock) {
           continue;
         }
@@ -143,7 +248,9 @@ export class PrivateEventStore {
         }
 
         events.push({
-          eventCommitmentIndex: entry.eventCommitmentIndex,
+          l2BlockNumber: entry.l2BlockNumber,
+          txIndexInBlock: entry.txIndexInBlock,
+          eventIndexInTx: entry.eventIndexInTx,
           event: {
             packedEvent: msgContent,
             l2BlockNumber: BlockNumber(entry.l2BlockNumber),
@@ -155,8 +262,63 @@ export class PrivateEventStore {
       }
     }
 
-    // Sort by eventCommitmentIndex only
-    events.sort((a, b) => a.eventCommitmentIndex - b.eventCommitmentIndex);
+    // Sort by block number, then by tx index within block, then by event index within tx
+    events.sort((a, b) => {
+      if (a.l2BlockNumber !== b.l2BlockNumber) {
+        return a.l2BlockNumber - b.l2BlockNumber;
+      }
+      if (a.txIndexInBlock !== b.txIndexInBlock) {
+        return a.txIndexInBlock - b.txIndexInBlock;
+      }
+      return a.eventIndexInTx - b.eventIndexInTx;
+    });
     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).
+   *
+   * We don't need staged writes for a rollback since it's handled in the context of a blockchain rewind.
+   *
+   * Rollbacks are handled by the BlockSynchronizer, which runs a DB transaction across stores when it detects a
+   * re-org, including setting the new anchor block after rolling back.
+   *
+   * So if anything fails in the process of rolling back any store, all DB changes occurring during rollbacks will be
+   * lost and the anchor block will not be updated; which means this code will eventually need to run again
+   * (i.e.: PXE will detect it's basing it work on an invalid block hash, then which re-triggers rewind).
+   *
+   * For further details, refer to `BlockSynchronizer#handleBlockStreamEvent`.
+   *
+   * IMPORTANT: This method must be called within a transaction to ensure atomicity.
+   */
+  public async rollback(blockNumber: number, synchedBlockNumber: number): Promise<void> {
+    let removedCount = 0;
+
+    for (let block = blockNumber + 1; block <= synchedBlockNumber; block++) {
+      const eventIds: string[] = [];
+      for await (const eventId of this.#eventsByBlockNumber.getValuesAsync(block)) {
+        eventIds.push(eventId);
+      }
+
+      if (eventIds.length > 0) {
+        await this.#eventsByBlockNumber.delete(block);
+
+        for (const eventId of eventIds) {
+          const entry = await this.#eventLogs.getAsync(eventId);
+          if (!entry) {
+            throw new Error(`Event log not found for eventId ${eventId}`);
+          }
+
+          await this.#eventLogs.delete(eventId);
+          await this.#seenLogs.delete(eventId);
+          await this.#eventsByContractScopeSelector.deleteValue(entry.lookupKey, eventId);
+
+          removedCount++;
+        }
+      }
+    }
+
+    this.logger.verbose(`Rolled back ${removedCount} private events after block ${blockNumber}`);
+  }
 }

package/src/storage/tagging_store/recipient_tagging_store.ts

@@ -1,6 +1,8 @@
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { DirectionalAppTaggingSecret } from '@aztec/stdlib/logs';
 
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
+
 /**
  * Data provider of tagging data used when syncing the logs as a recipient. The sender counterpart of this class
  * is called SenderTaggingStore. We have the providers separate for the sender and recipient because
@@ -8,46 +10,120 @@ import type { DirectionalAppTaggingSecret } from '@aztec/stdlib/logs';
  *
  * @dev Chain reorgs do not need to be handled here because both the finalized and aged indexes refer to finalized
  * blocks, which by definition cannot be affected by reorgs.
- *
- * TODO(benesjan): Relocate to yarn-project/pxe/src/storage/tagging_store
  */
-export class RecipientTaggingStore {
+export class RecipientTaggingStore implements StagedStore {
+  storeName: string = 'recipient_tagging';
+
   #store: AztecAsyncKVStore;
 
   #highestAgedIndex: AztecAsyncMap<string, number>;
   #highestFinalizedIndex: AztecAsyncMap<string, number>;
 
+  // jobId => secret => number
+  #highestAgedIndexForJob: Map<string, Map<string, number>>;
+
+  // jobId => secret => number
+  #highestFinalizedIndexForJob: Map<string, Map<string, number>>;
+
   constructor(store: AztecAsyncKVStore) {
     this.#store = store;
 
     this.#highestAgedIndex = this.#store.openMap('highest_aged_index');
     this.#highestFinalizedIndex = this.#store.openMap('highest_finalized_index');
+
+    this.#highestAgedIndexForJob = new Map();
+    this.#highestFinalizedIndexForJob = new Map();
+  }
+
+  #getHighestAgedIndexForJob(jobId: string): Map<string, number> {
+    let highestAgedIndexForJob = this.#highestAgedIndexForJob.get(jobId);
+    if (!highestAgedIndexForJob) {
+      highestAgedIndexForJob = new Map();
+      this.#highestAgedIndexForJob.set(jobId, highestAgedIndexForJob);
+    }
+    return highestAgedIndexForJob;
+  }
+
+  async #readHighestAgedIndex(jobId: string, secret: string): Promise<number | undefined> {
+    return this.#getHighestAgedIndexForJob(jobId).get(secret) ?? (await this.#highestAgedIndex.getAsync(secret));
+  }
+
+  #writeHighestAgedIndex(jobId: string, secret: string, index: number) {
+    this.#getHighestAgedIndexForJob(jobId).set(secret, index);
+  }
+
+  #getHighestFinalizedIndexForJob(jobId: string): Map<string, number> {
+    let jobStagedHighestFinalizedIndex = this.#highestFinalizedIndexForJob.get(jobId);
+    if (!jobStagedHighestFinalizedIndex) {
+      jobStagedHighestFinalizedIndex = new Map();
+      this.#highestFinalizedIndexForJob.set(jobId, jobStagedHighestFinalizedIndex);
+    }
+    return jobStagedHighestFinalizedIndex;
+  }
+
+  async #readHighestFinalizedIndex(jobId: string, secret: string): Promise<number | undefined> {
+    return (
+      this.#getHighestFinalizedIndexForJob(jobId).get(secret) ?? (await this.#highestFinalizedIndex.getAsync(secret))
+    );
+  }
+
+  #writeHighestFinalizedIndex(jobId: string, secret: string, index: number) {
+    this.#getHighestFinalizedIndexForJob(jobId).set(secret, index);
+  }
+
+  /**
+   * Writes all job-specific in-memory data to persistent storage.
+   *
+   * @remark This method must run in a DB transaction context. It's designed to be called from JobCoordinator#commitJob.
+   */
+  async commit(jobId: string): Promise<void> {
+    const highestAgedIndexForJob = this.#highestAgedIndexForJob.get(jobId);
+    if (highestAgedIndexForJob) {
+      for (const [secret, index] of highestAgedIndexForJob.entries()) {
+        await this.#highestAgedIndex.set(secret, index);
+      }
+    }
+
+    const highestFinalizedIndexForJob = this.#highestFinalizedIndexForJob.get(jobId);
+    if (highestFinalizedIndexForJob) {
+      for (const [secret, index] of highestFinalizedIndexForJob.entries()) {
+        await this.#highestFinalizedIndex.set(secret, index);
+      }
+    }
+
+    return this.discardStaged(jobId);
+  }
+
+  discardStaged(jobId: string): Promise<void> {
+    this.#highestAgedIndexForJob.delete(jobId);
+    this.#highestFinalizedIndexForJob.delete(jobId);
+    return Promise.resolve();
   }
 
-  getHighestAgedIndex(secret: DirectionalAppTaggingSecret): Promise<number | undefined> {
-    return this.#highestAgedIndex.getAsync(secret.toString());
+  getHighestAgedIndex(secret: DirectionalAppTaggingSecret, jobId: string): Promise<number | undefined> {
+    return this.#readHighestAgedIndex(jobId, secret.toString());
   }
 
-  async updateHighestAgedIndex(secret: DirectionalAppTaggingSecret, index: number): Promise<void> {
-    const currentIndex = await this.#highestAgedIndex.getAsync(secret.toString());
+  async updateHighestAgedIndex(secret: DirectionalAppTaggingSecret, index: number, jobId: string): Promise<void> {
+    const currentIndex = await this.#readHighestAgedIndex(jobId, secret.toString());
     if (currentIndex !== undefined && index <= currentIndex) {
       // Log sync should never set a lower highest aged index.
       throw new Error(`New highest aged index (${index}) must be higher than the current one (${currentIndex})`);
     }
-    await this.#highestAgedIndex.set(secret.toString(), index);
+    this.#writeHighestAgedIndex(jobId, secret.toString(), index);
   }
 
-  getHighestFinalizedIndex(secret: DirectionalAppTaggingSecret): Promise<number | undefined> {
-    return this.#highestFinalizedIndex.getAsync(secret.toString());
+  getHighestFinalizedIndex(secret: DirectionalAppTaggingSecret, jobId: string): Promise<number | undefined> {
+    return this.#readHighestFinalizedIndex(jobId, secret.toString());
   }
 
-  async updateHighestFinalizedIndex(secret: DirectionalAppTaggingSecret, index: number): Promise<void> {
-    const currentIndex = await this.#highestFinalizedIndex.getAsync(secret.toString());
+  async updateHighestFinalizedIndex(secret: DirectionalAppTaggingSecret, index: number, jobId: string): Promise<void> {
+    const currentIndex = await this.#readHighestFinalizedIndex(jobId, secret.toString());
     if (currentIndex !== undefined && index < currentIndex) {
       // Log sync should never set a lower highest finalized index but it can happen that it would try to set the same
       // one because we are loading logs from highest aged index + 1 and not from the highest finalized index.
       throw new Error(`New highest finalized index (${index}) must be higher than the current one (${currentIndex})`);
     }
-    await this.#highestFinalizedIndex.set(secret.toString(), index);
+    this.#writeHighestFinalizedIndex(jobId, secret.toString(), index);
   }
 }