refactor(pxe): migrate FactStore to BaseStagingStore (backport of aztec-labs-eng/aztec-node#98)

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

@@ -118,6 +118,7 @@ describe('Utility Execution test suite', () => {
     taggingSecretSourcesStore = mock<TaggingSecretSourcesStore>();
     capsuleStore = mock<CapsuleStore>();
     factStore = new FactStore(await openTmpStore('utility-exec-fact-test'));
+    factStore.beginChangeSet('test-change-set-id');
     privateEventStore = mock<PrivateEventStore>();
     contractSyncService = mock<ContractSyncService>();
     l2TipsStore = mock<L2TipsProvider>();

yarn-project/pxe/src/logs/log_service.test.ts

@@ -492,6 +492,7 @@ async function createTestLogService(
 ) {
   const keyStore = new KeyStore(await openTmpStore('test'));
   const recipientTaggingStore = new RecipientTaggingStore(await openTmpStore('test'));
+  recipientTaggingStore.beginChangeSet('test');
   const taggingSecretSourcesStore = new TaggingSecretSourcesStore(await openTmpStore('test'));
   const addressStore = new AddressStore(await openTmpStore('test'));
   const aztecNode = mock<AztecNode>();

yarn-project/pxe/src/operation_lifecycle.test.ts

@@ -136,6 +136,34 @@ describe('runOperation', () => {
     expect(notified).toEqual(['committed', 'discarded']);
   });
 
+  it('reports the operation error and still notifies contributors when the abort fails', async () => {
+    coordinator = new StagedWriteCoordinator({
+      kvStore: store,
+      stagedStores: [
+        {
+          storeName: 'undiscardable_store',
+          commitChangeSet: () => Promise.resolve(),
+          discardChangeSet: () => {
+            throw new Error('cannot discard');
+          },
+        },
+      ],
+    });
+    const notified: string[] = [];
+    const contributors: OperationContributor[] = [
+      {
+        onOperationEnd: (_, outcome) => {
+          notified.push(outcome);
+        },
+      },
+    ];
+
+    await expect(run(contributors, () => Promise.reject(new Error('operation failed'))).operation).rejects.toThrow(
+      'operation failed',
+    );
+    expect(notified).toEqual(['discarded']);
+  });
+
   /** Begins a change set and runs `fn` as an operation over it. */
   function run<T>(contributors: OperationContributor[], fn: () => Promise<T>) {
     const changeSetId = coordinator.begin();

yarn-project/pxe/src/operation_lifecycle.ts

@@ -52,7 +52,13 @@ export async function runOperation<T>(args: RunOperationArgs, fn: () => Promise<
   } catch (err) {
     log.verbose(`Aborting operation ${changeSetId}`, { changeSetId });
     await settleContributorsLoggingFailures(args);
-    stagedWriteCoordinator.abort(changeSetId);
+    try {
+      stagedWriteCoordinator.abort(changeSetId);
+    } catch (abortErr) {
+      // Nothing here can undo a failed abort, so it is logged, but the error that ended the operation is the one
+      // reported.
+      log.error(`Failed to abort operation ${changeSetId}`, abortErr, { changeSetId });
+    }
     notifyOperationEnd(args, 'discarded');
     throw err;
   }

yarn-project/pxe/src/storage/backwards_compatibility_tests/schema_tests.ts

@@ -223,6 +223,7 @@ export const SCHEMA_TESTS: readonly SchemaTest[] = [
     writeToStore: async kvStore => {
       const factStore = new FactStore(kvStore);
       const changeSetId = 'fixture-change-set';
+      factStore.beginChangeSet(changeSetId);
       const contract = AztecAddress.fromBigIntUnsafe(100n);
       const scope = AztecAddress.fromBigIntUnsafe(1n);
       const factCollectionTypeId = new Fr(7n);
@@ -536,6 +537,8 @@ export const SCHEMA_TESTS: readonly SchemaTest[] = [
       const recipientTaggingStore = new RecipientTaggingStore(kvStore);
 
       const changeSetId = 'fixture-change-set';
+      recipientTaggingStore.beginChangeSet(changeSetId);
+
       const secretA = new AppTaggingSecret(new Fr(2n), AztecAddress.fromBigIntUnsafe(3n));
       const secretB = new AppTaggingSecret(new Fr(5n), AztecAddress.fromBigIntUnsafe(7n));
       // A constrained secret keys under the `constrained:` prefix, so the snapshot pins both kinds side by side.

yarn-project/pxe/src/storage/base_staging_store.test.ts

@@ -15,7 +15,7 @@ describe('BaseStagingStore', () => {
     store = new TestStore(kv);
   });
 
-  describe('withChangeSet', () => {
+  describe('withChangeSetAndDb', () => {
     it('accepts operations between beginChangeSet and the end of the change set', async () => {
       store.beginChangeSet('cs1');
       await store.write('key', 1, 'cs1');
@@ -120,6 +120,41 @@ describe('BaseStagingStore', () => {
     });
   });
 
+  describe('withChangeSet', () => {
+    it('accepts operations between beginChangeSet and the end of the change set', async () => {
+      store.beginChangeSet('cs1');
+      await store.stage('key', 1, 'cs1');
+      await expect(store.staged('key', 'cs1')).resolves.toBe(1);
+    });
+
+    it('rejects staging for a change set that is not open', async () => {
+      await expect(store.stage('key', 1, 'cs1')).rejects.toThrow('Store "test": change set "cs1" is not open');
+    });
+
+    it('takes the same lock as withChangeSetAndDb, so it cannot interleave with an operation in flight', async () => {
+      store.beginChangeSet('cs1');
+      const gate = promiseWithResolvers<void>();
+      const order: string[] = [];
+
+      const inFlight = store.op(async changeSet => {
+        order.push('op-start');
+        await gate.promise;
+        changeSet.set('key', 1);
+        order.push('op-end');
+      }, 'cs1');
+      await tick();
+
+      const staged = store.stage('key', 2, 'cs1').then(() => order.push('stage-end'));
+      await tick();
+      expect(order).toEqual(['op-start']);
+
+      gate.resolve();
+      await Promise.all([inFlight, staged]);
+      expect(order).toEqual(['op-start', 'op-end', 'stage-end']);
+      await expect(store.readStaged('key', 'cs1')).resolves.toBe(2);
+    });
+  });
+
   describe('beginChangeSet', () => {
     it('rejects opening a change set while another is open', async () => {
       store.beginChangeSet('cs1');
@@ -181,6 +216,21 @@ describe('BaseStagingStore', () => {
   });
 
   describe('discardChangeSet', () => {
+    it('drops the staged writes without touching committed state', async () => {
+      store.beginChangeSet('cs1');
+      await store.write('key', 1, 'cs1');
+      await store.commitChangeSet('cs1');
+
+      store.beginChangeSet('cs2');
+      await store.write('key', 2, 'cs2');
+      store.discardChangeSet('cs2');
+
+      // Commit the next change set: a discarded change set's staged write must not ride along on it.
+      store.beginChangeSet('cs3');
+      await store.commitChangeSet('cs3');
+      await expect(store.committed('key')).resolves.toBe(1);
+    });
+
     it('leaves the open change set alone when discarding a different one', async () => {
       store.beginChangeSet('cs1');
       await store.write('key', 1, 'cs1');
@@ -263,19 +313,29 @@ class TestStore extends BaseStagingStore<Map<string, number>, TestDb> {
   }
 
   write(key: string, value: number, changeSetId: ChangeSetId): Promise<void> {
-    return this.withChangeSet(changeSetId, changeSet => {
+    return this.withChangeSetAndDb(changeSetId, changeSet => {
       changeSet.set(key, value);
       return Promise.resolve();
     });
   }
 
+  stage(key: string, value: number, changeSetId: ChangeSetId): Promise<void> {
+    return this.withChangeSet(changeSetId, changeSet => {
+      changeSet.set(key, value);
+    });
+  }
+
+  staged(key: string, changeSetId: ChangeSetId): Promise<number | undefined> {
+    return this.withChangeSet(changeSetId, changeSet => changeSet.get(key));
+  }
+
   // Runs an arbitrary operation body under the change set's lock.
   op<R>(fn: (changeSet: Map<string, number>) => Promise<R>, changeSetId: ChangeSetId): Promise<R> {
-    return this.withChangeSet(changeSetId, changeSet => fn(changeSet));
+    return this.withChangeSetAndDb(changeSetId, changeSet => fn(changeSet));
   }
 
   readStaged(key: string, changeSetId: ChangeSetId): Promise<number | undefined> {
-    return this.withChangeSet(changeSetId, changeSet => Promise.resolve(changeSet.get(key)));
+    return this.withChangeSetAndDb(changeSetId, changeSet => Promise.resolve(changeSet.get(key)));
   }
 
   committed(key: string): Promise<number | undefined> {

yarn-project/pxe/src/storage/base_staging_store.ts

@@ -12,12 +12,13 @@ import type { ChangeSetId, StagedStore } from './staged_write_coordinator.js';
  * abort) finds no matching change set and throws, so it cannot stage new data for a dead one.
  *
  * The staged data and the store's kv handles live in this base class and are only reachable through
- * {@link withChangeSet} (change set operations, with the DB read-only), {@link flushChangeSet} (the commit-time
- * write-back) and {@link applyRollback} (the reorg truncation).
+ * {@link withChangeSet} (change set operations, staged data only), {@link withChangeSetAndDb} (change set operations
+ * that also read the DB), {@link flushChangeSet} (the commit-time write-back) and {@link applyRollback} (the reorg
+ * truncation).
  *
- * The class is thread safe: an internal lock serializes the operations run through {@link withChangeSet}, so two of
- * them issued concurrently (e.g. under `Promise.all`) never interleave across awaits. Subclasses can therefore read
- * staged data and write it back without handling atomicity themselves.
+ * The class is thread safe: an internal lock serializes the operations run through {@link withChangeSet} and
+ * {@link withChangeSetAndDb}, so two of them issued concurrently (e.g. under `Promise.all`) never interleave across
+ * awaits. Subclasses can therefore read staged data and write it back without handling atomicity themselves.
  *
  * @typeParam TChangeSet - The in-memory buffer a change set accumulates its writes in, built empty by `buildChangeSet`
  * on every {@link beginChangeSet} and dropped when the change set ends. A store keyed by id might stage
@@ -117,40 +118,65 @@ export abstract class BaseStagingStore<TChangeSet, TDb> implements StagedStore,
 
   /**
    * Writes the change set's staged data to persistent storage. Runs inside the caller's transaction: it must not
-   * open a transaction of its own or call {@link withChangeSet}.
+   * open a transaction of its own or call {@link withChangeSetAndDb}.
    */
   protected abstract flushChangeSet(changeSet: TChangeSet, db: TDb): Promise<void>;
 
   /**
    * Deletes the state originating from blocks strictly above `toBlock`. Runs inside the transaction owned by
-   * {@link rollbackToBlock}'s caller: it must not open a transaction of its own or call {@link withChangeSet}.
+   * {@link rollbackToBlock}'s caller: it must not open a transaction of its own or call {@link withChangeSetAndDb}.
    */
   protected abstract applyRollback(toBlock: number, db: TDb): Promise<void>;
 
   /**
-   * Runs a change set operation (read or write). Takes the store's lock, opens a transaction, and calls `fn` with the
-   * change set's staged data and a read-only view of the DB (writes are staged in memory until {@link flushChangeSet}
-   * runs on commit).
+   * Runs a change set operation that reads the DB. Takes the store's lock, opens a transaction, and calls `fn` with
+   * the change set's staged data and a read-only view of the DB (writes are staged in memory until
+   * {@link flushChangeSet} runs on commit).
+   *
+   * Prefer {@link withChangeSet} unless `fn` actually reads the DB.
    *
    * The lock makes the store thread safe: two operations issued concurrently (e.g. under `Promise.all`) cannot
    * interleave across awaits, so `fn` can read staged data and write it back without handling atomicity itself.
    *
    * @throws If the change set is not open.
    */
-  protected async withChangeSet<R>(
+  protected withChangeSetAndDb<R>(
     changeSetId: ChangeSetId,
     fn: (changeSet: TChangeSet, db: ReadonlyDb<TDb>) => Promise<R>,
   ): Promise<R> {
+    return this.#runLocked(changeSetId, () =>
+      this.#store.transactionAsync(() => {
+        // Re-resolve after the wait: the change set may have been discarded while this operation queued on the lock
+        // or on the store's transaction queue.
+        const current = this.#currentOrThrow(changeSetId);
+        return fn(current.changeSet, this.#db);
+      }),
+    );
+  }
+
+  /**
+   * Runs a change set operation over the staged data alone: takes the store's lock and calls `fn` with the change
+   * set's staged data, opening no db transaction.
+   *
+   * Use {@link withChangeSetAndDb} instead when `fn` reads the DB.
+   *
+   * @throws If the change set is not open.
+   */
+  protected withChangeSet<R>(changeSetId: ChangeSetId, fn: (changeSet: TChangeSet) => R | Promise<R>): Promise<R> {
+    return this.#runLocked(changeSetId, async () => {
+      // Re-resolve after the wait: the change set may have been discarded while this operation queued on the lock.
+      const current = this.#currentOrThrow(changeSetId);
+      return await fn(current.changeSet);
+    });
+  }
+
+  async #runLocked<R>(changeSetId: ChangeSetId, run: () => Promise<R>): Promise<R> {
     const entered = this.#currentOrThrow(changeSetId);
     entered.inFlight++;
     try {
       await this.#lock.acquire();
       try {
-        return await this.#store.transactionAsync(() => {
-          // Re-resolve after the wait: the change set may have been discarded while this operation queued on the lock.
-          const current = this.#currentOrThrow(changeSetId);
-          return fn(current.changeSet, this.#db);
-        });
+        return await run();
       } finally {
         this.#lock.release();
       }

yarn-project/pxe/src/storage/fact_store/fact_service.test.ts

@@ -34,6 +34,7 @@ describe('FactService', () => {
   beforeEach(async () => {
     kv = await openTmpStore('fact-service-test');
     store = new FactStore(kv);
+    store.beginChangeSet(changeSetId);
   });
 
   it('delegates record+get for an allowed scope', async () => {