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/block_synchronizer/block_synchronizer.ts

@@ -46,6 +46,10 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
   public async handleBlockStreamEvent(event: L2BlockStreamEvent): Promise<void> {
     await this.l2TipsStore.handleBlockStreamEvent(event);
 
+    // NOTE: Sync operations write directly to main storage (no staging) because they update
+    // our local view to match chain state. This is idempotent - if a job fails mid-sync,
+    // the next sync will bring us back to the correct state.
+
     switch (event.type) {
       case 'blocks-added': {
         const lastBlock = event.blocks.at(-1)!.block;

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

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

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

@@ -0,0 +1,46 @@
+/**
+ * JobContext represents an active job in the PXE.
+ *
+ * During a job, all database writes go to a staging area (prefixed keys).
+ * Only on job success are they promoted to main storage. This provides
+ * job-level atomicity without requiring cross-store transactions.
+ */
+export class JobContext {
+  /** Unique identifier for this job */
+  readonly jobId: string;
+
+  /** Type of job (for logging/debugging) */
+  readonly jobType: string;
+
+  /** Prefix for staging keys: "job_{jobId}:" */
+  readonly stagingPrefix: string;
+
+  /** Timestamp when job started */
+  readonly startedAt: number;
+
+  constructor(jobId: string, jobType: string) {
+    this.jobId = jobId;
+    this.jobType = jobType;
+    this.stagingPrefix = `job_${jobId}:`;
+    this.startedAt = Date.now();
+  }
+
+  /**
+   * Generates a staging key from a main key.
+   * Prepends the staging prefix to the key.
+   */
+  stagingKey(mainKey: string): string {
+    return `${this.stagingPrefix}${mainKey}`;
+  }
+
+  /**
+   * Extracts the main key from a staging key.
+   * Removes the staging prefix from the key.
+   */
+  mainKey(stagingKey: string): string {
+    if (!stagingKey.startsWith(this.stagingPrefix)) {
+      throw new Error(`Key "${stagingKey}" does not have staging prefix "${this.stagingPrefix}"`);
+    }
+    return stagingKey.substring(this.stagingPrefix.length);
+  }
+}

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

@@ -0,0 +1,135 @@
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+import { openTmpStore } from '@aztec/kv-store/lmdb-v2';
+
+import { jest } from '@jest/globals';
+
+import { JobContext } from './job_context.js';
+import { JobCoordinator, type StagedStore } from './job_coordinator.js';
+
+describe('JobCoordinator', () => {
+  let store: AztecAsyncKVStore;
+  let coordinator: JobCoordinator;
+
+  beforeEach(async () => {
+    store = await openTmpStore('job_coordinator_test');
+    coordinator = new JobCoordinator(store);
+  });
+
+  describe('beginJob', () => {
+    it('creates a new job context', () => {
+      const context = coordinator.beginJob('test_job');
+
+      expect(context).toBeInstanceOf(JobContext);
+      expect(context.jobType).toBe('test_job');
+      expect(context.jobId).toBeDefined();
+      expect(context.stagingPrefix).toContain('job_');
+    });
+
+    // Note: we could eventually be relax this if we want more concurrency,
+    // but it's good to start with this guardrail
+    it('throws if job already in progress', () => {
+      coordinator.beginJob('first_job');
+      expect(() => coordinator.beginJob('second_job')).toThrow(/already in progress/);
+    });
+
+    it('tracks job in progress', () => {
+      coordinator.beginJob('test_job');
+      expect(coordinator.hasJobInProgress()).toBe(true);
+    });
+  });
+
+  describe('commitJob', () => {
+    it('clears job marker on commit', async () => {
+      const context = coordinator.beginJob('test_job');
+      await coordinator.commitJob(context);
+      expect(coordinator.hasJobInProgress()).toBe(false);
+    });
+
+    it('throws if no matching job in progress', async () => {
+      const context = coordinator.beginJob('test_job');
+      await coordinator.commitJob(context);
+      await expect(coordinator.commitJob(context)).rejects.toThrow(/no matching job/);
+    });
+
+    it('calls commit on registered stores', async () => {
+      const commitMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const discardStagedMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const mockStore: StagedStore = {
+        storeName: 'mock_store',
+        commit: commitMock,
+        discardStaged: discardStagedMock,
+      };
+
+      coordinator.registerStore(mockStore);
+
+      const context = coordinator.beginJob('test_job');
+
+      await coordinator.commitJob(context);
+
+      expect(commitMock).toHaveBeenCalledWith(context);
+    });
+  });
+
+  describe('abortJob', () => {
+    it('clears job marker on abort', async () => {
+      const context = coordinator.beginJob('test_job');
+
+      await coordinator.abortJob(context);
+
+      expect(coordinator.hasJobInProgress()).toBe(false);
+    });
+
+    it('calls discardStaged on all registered stores', async () => {
+      const commitMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const discardStagedMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const mockStore: StagedStore = {
+        storeName: 'mock_store',
+        commit: commitMock,
+        discardStaged: discardStagedMock,
+      };
+
+      coordinator.registerStore(mockStore);
+
+      const context = coordinator.beginJob('test_job');
+
+      await coordinator.abortJob(context);
+
+      expect(discardStagedMock).toHaveBeenCalledWith(context.stagingPrefix);
+    });
+  });
+
+  describe('registerStore', () => {
+    it('throws on duplicate registration', () => {
+      const commitMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const discardStagedMock = jest.fn<() => Promise<void>>().mockResolvedValue(undefined);
+      const mockStore: StagedStore = {
+        storeName: 'mock_store',
+        commit: commitMock,
+        discardStaged: discardStagedMock,
+      };
+
+      coordinator.registerStore(mockStore);
+
+      expect(() => coordinator.registerStore(mockStore)).toThrow(/already registered/);
+    });
+  });
+});
+
+describe('JobContext', () => {
+  it('generates staging keys', () => {
+    const context = new JobContext('abc123', 'test');
+    expect(context.stagingKey('notes')).toBe('job_abc123:notes');
+    expect(context.stagingKey('data:key')).toBe('job_abc123:data:key');
+  });
+
+  it('extracts main keys from staging keys', () => {
+    const context = new JobContext('abc123', 'test');
+    expect(context.mainKey('job_abc123:notes')).toBe('notes');
+    expect(context.mainKey('job_abc123:data:key')).toBe('data:key');
+  });
+
+  it('throws when extracting main key from non-staging key', () => {
+    const context = new JobContext('abc123', 'test');
+    expect(() => context.mainKey('notes')).toThrow(/does not have staging prefix/);
+  });
+});

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

@@ -0,0 +1,160 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+
+import { JobContext } from './job_context.js';
+
+/**
+ * Interface that data providers must implement to support staged writes.
+ */
+export interface StagedStore {
+  /** Unique name identifying this provider (used for tracking affected stores) */
+  readonly storeName: string;
+
+  /**
+   * Commits staged data to main storage.
+   * Should be called within a transaction for atomicity.
+   *
+   * @param context - The job context containing the staging prefix
+   */
+  commit(context: JobContext): Promise<void>;
+
+  /**
+   * Discards staged data without committing.
+   * Called on abort or during recovery.
+   *
+   * @param stagingPrefix - The prefix used for staging keys
+   */
+  discardStaged(stagingPrefix: string): Promise<void>;
+}
+
+/**
+ * JobCoordinator manages job lifecycle and provides crash resilience for PXE operations.
+ *
+ * It uses a staged writes pattern:
+ * 1. When a job begins, a context is created with a unique staging prefix
+ * 2. During the job, all writes go to staging (prefixed keys)
+ * 3. On commit, staging is promoted to main storage
+ * 4. On abort, staged data is discarded
+ *
+ * Note: PXE should only rely on a single JobCoordinator instance, so it can eventually
+ * orchestrate concurrent jobs. Right now it doesn't make a difference because we're
+ * using a job queue with concurrency=1.
+ */
+export class JobCoordinator {
+  private readonly log = createLogger('pxe:job_coordinator');
+
+  /** The underlying KV store */
+  kvStore: AztecAsyncKVStore;
+
+  /** Current job context (in-memory only) */
+  #currentJob: JobContext | undefined;
+
+  /** All registered staged stores */
+  #stores: Map<string, StagedStore> = new Map();
+
+  constructor(kvStore: AztecAsyncKVStore) {
+    this.kvStore = kvStore;
+  }
+
+  /**
+   * Registers a staged store.
+   * Must be called during initialization for all stores that need staging support.
+   */
+  registerStore(store: StagedStore): void {
+    if (this.#stores.has(store.storeName)) {
+      throw new Error(`Store "${store.storeName}" is already registered`);
+    }
+    this.#stores.set(store.storeName, store);
+    this.log.debug(`Registered staged store: ${store.storeName}`);
+  }
+
+  /**
+   * Registers multiple staged stores.
+   */
+  registerStores(stores: StagedStore[]): void {
+    for (const store of stores) {
+      this.registerStore(store);
+    }
+  }
+
+  /**
+   * Begins a new job and returns a context for staged writes.
+   *
+   * @param jobType - Type of job (for logging/debugging)
+   * @returns JobContext to pass to write operations
+   */
+  beginJob(jobType: string): JobContext {
+    if (this.#currentJob) {
+      throw new Error(
+        `Cannot begin job "${jobType}": job "${this.#currentJob.jobType}" (${this.#currentJob.jobId}) is already in progress. ` +
+          `This should not happen - ensure jobs are properly committed or aborted.`,
+      );
+    }
+
+    const jobId = randomBytes(8).toString('hex');
+    const context = new JobContext(jobId, jobType);
+    this.#currentJob = context;
+
+    this.log.debug(`Started job ${jobId} (type: ${jobType})`);
+    return context;
+  }
+
+  /**
+   * Commits a job by promoting all staged data to main storage.
+   *
+   * @param context - The job context returned from beginJob
+   */
+  async commitJob(context: JobContext): Promise<void> {
+    if (!this.#currentJob || this.#currentJob.jobId !== context.jobId) {
+      throw new Error(
+        `Cannot commit job ${context.jobId}: no matching job in progress. ` +
+          `Current job: ${this.#currentJob ? this.#currentJob.jobId : 'none'}`,
+      );
+    }
+
+    this.log.debug(`Committing job ${context.jobId}`);
+
+    // Commit all stores atomically in a single transaction.
+    // Each store's commit is a no-op if it has no staged data (but that's up to each store to handle).
+    await this.kvStore.transactionAsync(async () => {
+      for (const store of this.#stores.values()) {
+        await store.commit(context);
+      }
+    });
+
+    this.#currentJob = undefined;
+    this.log.debug(`Job ${context.jobId} committed successfully`);
+  }
+
+  /**
+   * Aborts a job by discarding all staged data.
+   *
+   * @param context - The job context returned from beginJob
+   */
+  async abortJob(context: JobContext): Promise<void> {
+    if (!this.#currentJob || this.#currentJob.jobId !== context.jobId) {
+      // Job may have already been aborted or never started properly
+      this.log.warn(`Abort called for job ${context.jobId} but current job is ${this.#currentJob?.jobId ?? 'none'}`);
+    }
+
+    this.log.debug(`Aborting job ${context.jobId}`);
+
+    // Discard staging atomically
+    await this.kvStore.transactionAsync(async () => {
+      for (const store of this.#stores.values()) {
+        await store.discardStaged(context.stagingPrefix);
+      }
+    });
+
+    this.#currentJob = undefined;
+    this.log.debug(`Job ${context.jobId} aborted`);
+  }
+
+  /**
+   * Checks if there's a job currently in progress.
+   */
+  hasJobInProgress(): boolean {
+    return this.#currentJob !== undefined;
+  }
+}