@aztec/pxe 0.0.1-commit.96bb3f7 → 0.0.1-commit.a072138

package/src/storage/address_store/address_store.ts

@@ -43,23 +43,23 @@ export class AddressStore {
     });
   }
 
-  async #getCompleteAddress(address: AztecAddress): Promise<CompleteAddress | undefined> {
-    const index = await this.#completeAddressIndex.getAsync(address.toString());
-    if (index === undefined) {
-      return undefined;
-    }
-
-    const value = await this.#completeAddresses.atAsync(index);
-    return value ? await CompleteAddress.fromBuffer(value) : undefined;
-  }
-
   getCompleteAddress(account: AztecAddress): Promise<CompleteAddress | undefined> {
-    return this.#getCompleteAddress(account);
+    return this.#store.transactionAsync(async () => {
+      const index = await this.#completeAddressIndex.getAsync(account.toString());
+      if (index === undefined) {
+        return undefined;
+      }
+
+      const value = await this.#completeAddresses.atAsync(index);
+      return value ? await CompleteAddress.fromBuffer(value) : undefined;
+    });
   }
 
-  async getCompleteAddresses(): Promise<CompleteAddress[]> {
-    return await Promise.all(
-      (await toArray(this.#completeAddresses.valuesAsync())).map(v => CompleteAddress.fromBuffer(v)),
-    );
+  getCompleteAddresses(): Promise<CompleteAddress[]> {
+    return this.#store.transactionAsync(async () => {
+      return await Promise.all(
+        (await toArray(this.#completeAddresses.valuesAsync())).map(v => CompleteAddress.fromBuffer(v)),
+      );
+    });
   }
 }

package/src/storage/anchor_block_store/anchor_block_store.ts

@@ -10,6 +10,14 @@ export class AnchorBlockStore {
     this.#synchronizedHeader = this.#store.openSingleton('header');
   }
 
+  /**
+   * Sets the currently synchronized block
+   *
+   * Important: this method is only called from BlockSynchronizer, and since we need it to run atomically with other
+   * stores in the case of a reorg, it MUST NOT be wrapped in a `transactionAsync` call. Doing so would result in a
+   * deadlock when the backend is IndexedDB, because `transactionAsync` is not designed to support reentrancy.
+   *
+   */
   async setHeader(header: BlockHeader): Promise<void> {
     await this.#synchronizedHeader.set(header.toBuffer());
   }

package/src/storage/capsule_store/capsule_store.ts

@@ -3,12 +3,21 @@ import { type Logger, createLogger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 
-export class CapsuleStore {
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
+
+export class CapsuleStore implements StagedStore {
+  readonly storeName = 'capsule';
+
   #store: AztecAsyncKVStore;
 
   // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${key}`
   #capsules: AztecAsyncMap<string, Buffer>;
 
+  // jobId => `${contractAddress}:${key}` => capsule data
+  // when `#stagedCapsules.get('some-job-id').get('${some-contract-address:some-key') === null`,
+  // it signals that the capsule was deleted during the job, so it needs to be deleted on commit
+  #stagedCapsules: Map<string, Map<string, Buffer | null>>;
+
   logger: Logger;
 
   constructor(store: AztecAsyncKVStore) {
@@ -16,21 +25,120 @@ export class CapsuleStore {
 
     this.#capsules = this.#store.openMap('capsules');
 
+    this.#stagedCapsules = new Map();
+
     this.logger = createLogger('pxe:capsule-data-provider');
   }
 
+  /**
+   * Given a job denoted by `jobId`, it returns the
+   * capsules that said job has interacted with.
+   *
+   * Capsules that haven't been committed to persistence KV storage
+   * are kept in-memory in `#stagedCapsules`, this method provides a convenient
+   * way to access that in-memory collection of data.
+   *
+   * @param jobId
+   * @returns
+   */
+  #getJobStagedCapsules(jobId: string): Map<string, Buffer | null> {
+    let jobStagedCapsules = this.#stagedCapsules.get(jobId);
+    if (!jobStagedCapsules) {
+      jobStagedCapsules = new Map();
+      this.#stagedCapsules.set(jobId, jobStagedCapsules);
+    }
+    return jobStagedCapsules;
+  }
+
+  /**
+   * Reads a capsule's slot from the staged version of the data associated to the given jobId.
+   *
+   * If it is not there, it reads it from the KV store.
+   */
+  async #getFromStage(jobId: string, dbSlotKey: string): Promise<Buffer | null | undefined> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+    const staged: Buffer | null | undefined = jobStagedCapsules.get(dbSlotKey);
+
+    // Always issue DB read to keep IndexedDB transaction alive, even if the value is in the job staged data. This
+    // keeps IndexedDB transactions alive (they auto-commit when a new micro-task starts and there are no pending read
+    // requests). The staged value still takes precedence if it exists (including null for deletions).
+    const dbValue = await this.#loadCapsuleFromDb(dbSlotKey);
+
+    return staged !== undefined ? staged : dbValue;
+  }
+
+  /**
+   * Writes a capsule to the stage of a job.
+   */
+  #setOnStage(jobId: string, dbSlotKey: string, capsuleData: Buffer) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, capsuleData);
+  }
+
+  /**
+   * Deletes a capsule on the stage of a job. Note the capsule will still
+   * exist in storage until the job is committed.
+   */
+  #deleteOnStage(jobId: string, dbSlotKey: string) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, null);
+  }
+
+  async #loadCapsuleFromDb(dbSlotKey: string): Promise<Buffer | null> {
+    const dataBuffer = await this.#capsules.getAsync(dbSlotKey);
+    if (!dataBuffer) {
+      return null;
+    }
+
+    return dataBuffer;
+  }
+
+  /**
+   * Commits staged data to main storage.
+   * Called by JobCoordinator when a job completes successfully.
+   * Note: JobCoordinator wraps all commits in a single transaction, so we don't
+   * need our own transactionAsync here (and using one would deadlock on IndexedDB).
+   * @param jobId - The jobId identifying which staged data to commit
+   */
+  async commit(jobId: string): Promise<void> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+
+    for (const [key, value] of jobStagedCapsules) {
+      // In the write stage, we represent deleted capsules with null
+      // (as opposed to undefined, which denotes there was never a capsule there to begin with).
+      // So we delete from actual KV store here.
+      if (value === null) {
+        await this.#capsules.delete(key);
+      } else {
+        await this.#capsules.set(key, value);
+      }
+    }
+
+    this.#stagedCapsules.delete(jobId);
+  }
+
+  /**
+   * Discards staged data without committing.
+   */
+  discardStaged(jobId: string): Promise<void> {
+    this.#stagedCapsules.delete(jobId);
+    return Promise.resolve();
+  }
+
   /**
    * Stores arbitrary information in a per-contract non-volatile database, which can later be retrieved with `loadCapsule`.
    * * If data was already stored at this slot, it is overwritten.
    * @param contractAddress - The contract address to scope the data under.
    * @param slot - The slot in the database in which to store the value. Slots need not be contiguous.
    * @param capsule - An array of field elements representing the capsule.
+   * @param jobId - The context in which this store will be visible until PXE decides to persist it to underlying KV store
    * @remarks A capsule is a "blob" of data that is passed to the contract through an oracle. It works similarly
    * to public contract storage in that it's indexed by the contract address and storage slot but instead of the global
    * network state it's backed by local PXE db.
    */
-  async storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[]): Promise<void> {
-    await this.#capsules.set(dbSlotToKey(contractAddress, slot), Buffer.concat(capsule.map(value => value.toBuffer())));
+  storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[], jobId: string) {
+    const dbSlotKey = dbSlotToKey(contractAddress, slot);
+
+    // A store overrides any pre-existing data on the slot
+    this.#setOnStage(jobId, dbSlotKey, Buffer.concat(capsule.map(value => value.toBuffer())));
   }
 
   /**
@@ -39,8 +147,8 @@ export class CapsuleStore {
    * @param slot - The slot in the database to read.
    * @returns The stored data or `null` if no data is stored under the slot.
    */
-  async loadCapsule(contractAddress: AztecAddress, slot: Fr): Promise<Fr[] | null> {
-    const dataBuffer = await this.#capsules.getAsync(dbSlotToKey(contractAddress, slot));
+  async loadCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string): Promise<Fr[] | null> {
+    const dataBuffer = await this.#getFromStage(jobId, dbSlotToKey(contractAddress, slot));
     if (!dataBuffer) {
       this.logger.trace(`Data not found for contract ${contractAddress.toString()} and slot ${slot.toString()}`);
       return null;
@@ -57,8 +165,9 @@ export class CapsuleStore {
    * @param contractAddress - The contract address under which the data is scoped.
    * @param slot - The slot in the database to delete.
    */
-  async deleteCapsule(contractAddress: AztecAddress, slot: Fr): Promise<void> {
-    await this.#capsules.delete(dbSlotToKey(contractAddress, slot));
+  deleteCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string) {
+    // When we commit this, we will interpret null as a deletion, so we'll propagate the delete to the KV store
+    this.#deleteOnStage(jobId, dbSlotToKey(contractAddress, slot));
   }
 
   /**
@@ -72,13 +181,22 @@ export class CapsuleStore {
    * @param dstSlot - The first slot to copy to.
    * @param numEntries - The number of entries to copy.
    */
-  copyCapsule(contractAddress: AztecAddress, srcSlot: Fr, dstSlot: Fr, numEntries: number): Promise<void> {
+  copyCapsule(
+    contractAddress: AztecAddress,
+    srcSlot: Fr,
+    dstSlot: Fr,
+    numEntries: number,
+    jobId: string,
+  ): Promise<void> {
+    // This transactional context gives us "copy atomicity":
+    // there shouldn't be concurrent writes to what's being copied here.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // In order to support overlapping source and destination regions, we need to check the relative positions of source
       // and destination. If destination is ahead of source, then by the time we overwrite source elements using forward
       // indexes we'll have already read those. On the contrary, if source is ahead of destination we need to use backward
       // indexes to avoid reading elements that've been overwritten.
-
       const indexes = Array.from(Array(numEntries).keys());
       if (srcSlot.lt(dstSlot)) {
         indexes.reverse();
@@ -88,12 +206,12 @@ export class CapsuleStore {
         const currentSrcSlot = dbSlotToKey(contractAddress, srcSlot.add(new Fr(i)));
         const currentDstSlot = dbSlotToKey(contractAddress, dstSlot.add(new Fr(i)));
 
-        const toCopy = await this.#capsules.getAsync(currentSrcSlot);
+        const toCopy = await this.#getFromStage(jobId, currentSrcSlot);
         if (!toCopy) {
           throw new Error(`Attempted to copy empty slot ${currentSrcSlot} for contract ${contractAddress.toString()}`);
         }
 
-        await this.#capsules.set(currentDstSlot, toCopy);
+        this.#setOnStage(jobId, currentDstSlot, toCopy);
       }
     });
   }
@@ -106,35 +224,45 @@ export class CapsuleStore {
    * @param baseSlot - The slot where the array length is stored
    * @param content - Array of capsule data to append
    */
-  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]): Promise<void> {
+  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string): Promise<void> {
+    // We wrap this in a transaction to serialize concurrent calls from Promise.all.
+    // Without this, concurrent appends to the same array could race: both read length=0,
+    // both write at the same slots, one overwrites the other.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const lengthData = await this.loadCapsule(contractAddress, baseSlot);
+      const lengthData = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const currentLength = lengthData ? lengthData[0].toNumber() : 0;
 
       // Store each capsule at consecutive slots after baseSlot + 1 + currentLength
       for (let i = 0; i < content.length; i++) {
         const nextSlot = arraySlot(baseSlot, currentLength + i);
-        await this.storeCapsule(contractAddress, nextSlot, content[i]);
+        this.storeCapsule(contractAddress, nextSlot, content[i], jobId);
       }
 
       // Update length to include all new capsules
       const newLength = currentLength + content.length;
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)]);
+      this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)], jobId);
     });
   }
 
-  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr): Promise<Fr[][]> {
+  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, jobId: string): Promise<Fr[][]> {
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "read array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const length = maybeLength ? maybeLength[0].toBigInt() : 0n;
 
       const values: Fr[][] = [];
 
       // Read each capsule at consecutive slots after baseSlot
       for (let i = 0; i < length; i++) {
-        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i));
+        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i), jobId);
         if (currentValue == undefined) {
           throw new Error(
             `Expected non-empty value at capsule array in base slot ${baseSlot} at index ${i} for contract ${contractAddress}`,
@@ -148,23 +276,31 @@ export class CapsuleStore {
     });
   }
 
-  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]) {
+  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string) {
+    // This transactional context in theory isn't so critical now because we aren't
+    // writing to DB so if there's exceptions midway and it blows up, no visible impact
+    // to persistent storage will happen.
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "write array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.loadCapsule(contractAddress, baseSlot, jobId);
       const originalLength = maybeLength ? maybeLength[0].toNumber() : 0;
 
       // Set the new length
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)]);
+      this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)], jobId);
 
       // Store the new content, possibly overwriting existing values
       for (let i = 0; i < content.length; i++) {
-        await this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i]);
+        this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i], jobId);
       }
 
       // Clear any stragglers
       for (let i = content.length; i < originalLength; i++) {
-        await this.deleteCapsule(contractAddress, arraySlot(baseSlot, i));
+        this.deleteCapsule(contractAddress, arraySlot(baseSlot, i), jobId);
       }
     });
   }

package/src/storage/contract_store/contract_store.ts

@@ -3,7 +3,6 @@ import type { Fr } from '@aztec/foundation/curves/bn254';
 import { toArray } from '@aztec/foundation/iterable';
 import type { MembershipWitness } from '@aztec/foundation/trees';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
-import { isProtocolContract } from '@aztec/protocol-contracts';
 import {
   type ContractArtifact,
   type FunctionAbi,
@@ -43,10 +42,12 @@ export class ContractStore {
   /** Map from contract address to contract class id */
   #contractClassIdMap: Map<string, Fr> = new Map();
 
+  #store: AztecAsyncKVStore;
   #contractArtifacts: AztecAsyncMap<string, Buffer>;
   #contractInstances: AztecAsyncMap<string, Buffer>;
 
   constructor(store: AztecAsyncKVStore) {
+    this.#store = store;
     this.#contractArtifacts = store.openMap('contract_artifacts');
     this.#contractInstances = store.openMap('contracts_instances');
   }
@@ -54,6 +55,7 @@ export class ContractStore {
   // Setters
 
   public async addContractArtifact(id: Fr, contract: ContractArtifact): Promise<void> {
+    // Validation outside transactionAsync - these are not DB operations
     const privateFunctions = contract.functions.filter(
       functionArtifact => functionArtifact.functionType === FunctionType.PRIVATE,
     );
@@ -70,7 +72,9 @@ export class ContractStore {
       throw new Error('Repeated function selectors of private functions');
     }
 
-    await this.#contractArtifacts.set(id.toString(), contractArtifactToBuffer(contract));
+    await this.#store.transactionAsync(() =>
+      this.#contractArtifacts.set(id.toString(), contractArtifactToBuffer(contract)),
+    );
   }
 
   async addContractInstance(contract: ContractInstanceWithAddress): Promise<void> {
@@ -124,21 +128,27 @@ export class ContractStore {
 
   // Public getters
 
-  async getContractsAddresses(): Promise<AztecAddress[]> {
-    const keys = await toArray(this.#contractInstances.keysAsync());
-    return keys.map(AztecAddress.fromString);
+  getContractsAddresses(): Promise<AztecAddress[]> {
+    return this.#store.transactionAsync(async () => {
+      const keys = await toArray(this.#contractInstances.keysAsync());
+      return keys.map(AztecAddress.fromString);
+    });
   }
 
   /** Returns a contract instance for a given address. Throws if not found. */
-  public async getContractInstance(contractAddress: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
-    const contract = await this.#contractInstances.getAsync(contractAddress.toString());
-    return contract && SerializableContractInstance.fromBuffer(contract).withAddress(contractAddress);
+  public getContractInstance(contractAddress: AztecAddress): Promise<ContractInstanceWithAddress | undefined> {
+    return this.#store.transactionAsync(async () => {
+      const contract = await this.#contractInstances.getAsync(contractAddress.toString());
+      return contract && SerializableContractInstance.fromBuffer(contract).withAddress(contractAddress);
+    });
   }
 
-  public async getContractArtifact(contractClassId: Fr): Promise<ContractArtifact | undefined> {
-    const contract = await this.#contractArtifacts.getAsync(contractClassId.toString());
-    // TODO(@spalladino): AztecAsyncMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
-    return contract && contractArtifactFromBuffer(Buffer.from(contract));
+  public getContractArtifact(contractClassId: Fr): Promise<ContractArtifact | undefined> {
+    return this.#store.transactionAsync(async () => {
+      const contract = await this.#contractArtifacts.getAsync(contractClassId.toString());
+      // TODO(@spalladino): AztecAsyncMap lies and returns Uint8Arrays instead of Buffers, hence the extra Buffer.from.
+      return contract && contractArtifactFromBuffer(Buffer.from(contract));
+    });
   }
 
   /** Returns a contract class for a given class id. Throws if not found. */
@@ -317,23 +327,4 @@ export class ContractStore {
       returnTypes: functionDao.returnTypes,
     };
   }
-
-  // Synchronize target contract data
-  public async syncPrivateState(
-    contractAddress: AztecAddress,
-    functionToInvokeAfterSync: FunctionSelector | null,
-    utilityExecutor: (privateSyncCall: FunctionCall) => Promise<any>,
-  ) {
-    // Protocol contracts don't have private state to sync
-    if (!isProtocolContract(contractAddress)) {
-      const syncPrivateStateFunctionCall = await this.getFunctionCall('sync_private_state', [], contractAddress);
-      if (functionToInvokeAfterSync && functionToInvokeAfterSync.equals(syncPrivateStateFunctionCall.selector)) {
-        throw new Error(
-          'Forbidden `sync_private_state` invocation. `sync_private_state` can only be invoked by PXE, manual execution can lead to inconsistencies.',
-        );
-      }
-
-      return utilityExecutor(syncPrivateStateFunctionCall);
-    }
-  }
 }

package/src/storage/metadata.ts

@@ -1 +1 @@
-export const PXE_DATA_SCHEMA_VERSION = 2;
+export const PXE_DATA_SCHEMA_VERSION = 3;