@aztec/prover-client 0.0.1-commit.f1df4d2 → 0.0.1-commit.f224bb98b

package/src/orchestrator/orchestrator.ts

@@ -12,6 +12,7 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { AbortError } from '@aztec/foundation/error';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
+import { SerialQueue } from '@aztec/foundation/queue';
 import { assertLength } from '@aztec/foundation/serialize';
 import { pushTestData } from '@aztec/foundation/testing';
 import { elapsed } from '@aztec/foundation/timer';
@@ -71,11 +72,6 @@ import { EpochProvingState, type ProvingResult, type TreeSnapshots } from './epo
 import { ProvingOrchestratorMetrics } from './orchestrator_metrics.js';
 import { TxProvingState } from './tx-proving-state.js';
 
-type WorldStateFork = {
-  fork: MerkleTreeWriteOperations;
-  cleanupPromise: Promise<void> | undefined;
-};
-
 /**
  * Implements an event driven proving scheduler to build the recursive proof tree. The idea being:
  * 1. Transactions are provided to the scheduler post simulation.
@@ -97,19 +93,22 @@ export class ProvingOrchestrator implements EpochProver {
   private provingPromise: Promise<ProvingResult> | undefined = undefined;
   private metrics: ProvingOrchestratorMetrics;
   // eslint-disable-next-line aztec-custom/no-non-primitive-in-collections
-  private dbs: Map<BlockNumber, WorldStateFork> = new Map();
+  private dbs: Map<BlockNumber, MerkleTreeWriteOperations> = new Map();
   private logger: Logger;
+  private deferredJobQueue = new SerialQueue();
 
   constructor(
     private dbProvider: ReadonlyWorldStateAccess & ForkMerkleTreeOperations,
     private prover: ServerCircuitProver,
     private readonly proverId: EthAddress,
     private readonly cancelJobsOnStop: boolean = false,
+    private readonly enqueueConcurrency: number,
     telemetryClient: TelemetryClient = getTelemetryClient(),
     bindings?: LoggerBindings,
   ) {
     this.logger = createLogger('prover-client:orchestrator', bindings);
     this.metrics = new ProvingOrchestratorMetrics(telemetryClient, 'ProvingOrchestrator');
+    this.deferredJobQueue.start(this.enqueueConcurrency);
   }
 
   get tracer(): Tracer {
@@ -124,9 +123,11 @@ export class ProvingOrchestrator implements EpochProver {
     return this.dbs.size;
   }
 
-  public stop(): Promise<void> {
+  public async stop(): Promise<void> {
+    // Grab the old queue before cancel() replaces it, so we can await its draining.
+    const oldQueue = this.deferredJobQueue;
     this.cancel();
-    return Promise.resolve();
+    await oldQueue.cancel();
   }
 
   public startNewEpoch(
@@ -182,7 +183,7 @@ export class ProvingOrchestrator implements EpochProver {
     const db = await this.dbProvider.fork(lastBlockNumber);
 
     const firstBlockNumber = BlockNumber(lastBlockNumber + 1);
-    this.dbs.set(firstBlockNumber, { fork: db, cleanupPromise: undefined });
+    this.dbs.set(firstBlockNumber, db);
 
     // Get archive sibling path before any block in this checkpoint lands.
     const lastArchiveSiblingPath = await getLastSiblingPath(MerkleTreeId.ARCHIVE, db);
@@ -240,9 +241,9 @@ export class ProvingOrchestrator implements EpochProver {
     if (!this.dbs.has(blockNumber)) {
       // Fork world state at the end of the immediately previous block
       const db = await this.dbProvider.fork(BlockNumber(blockNumber - 1));
-      this.dbs.set(blockNumber, { fork: db, cleanupPromise: undefined });
+      this.dbs.set(blockNumber, db);
     }
-    const db = this.dbs.get(blockNumber)!.fork;
+    const db = this.getDbForBlock(blockNumber);
 
     // Get archive snapshot and sibling path before any txs in this block lands.
     const lastArchiveTreeSnapshot = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
@@ -317,7 +318,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     this.logger.info(`Adding ${txs.length} transactions to block ${blockNumber}`);
 
-    const db = this.dbs.get(blockNumber)!.fork;
+    const db = this.getDbForBlock(blockNumber);
     const lastArchive = provingState.lastArchiveTreeSnapshot;
     const newL1ToL2MessageTreeSnapshot = provingState.newL1ToL2MessageTreeSnapshot;
     const spongeBlobState = provingState.getStartSpongeBlob().clone();
@@ -445,14 +446,20 @@ export class ProvingOrchestrator implements EpochProver {
       throw new Error('Block header mismatch');
     }
 
-    // Get db for this block
-    const db = this.dbs.get(provingState.blockNumber)!.fork;
+    // Get db for this block and remove from map — no other code should use it after this point.
+    const db = this.getDbForBlock(provingState.blockNumber);
+    this.dbs.delete(provingState.blockNumber);
 
-    // Update the archive tree, so we're ready to start processing the next block:
-    this.logger.verbose(
-      `Updating archive tree with block ${provingState.blockNumber} header ${(await header.hash()).toString()}`,
-    );
-    await db.updateArchive(header);
+    // Update the archive tree, capture the snapshot, and close the fork deterministically.
+    try {
+      this.logger.verbose(
+        `Updating archive tree with block ${provingState.blockNumber} header ${(await header.hash()).toString()}`,
+      );
+      await db.updateArchive(header);
+      provingState.setBuiltArchive(await getTreeSnapshot(MerkleTreeId.ARCHIVE, db));
+    } finally {
+      await db.close();
+    }
 
     await this.verifyBuiltBlockAgainstSyncedState(provingState);
 
@@ -472,6 +479,13 @@ export class ProvingOrchestrator implements EpochProver {
       this.logger.debug('Block root rollup proof not built yet, skipping header check.');
       return;
     }
+
+    const newArchive = provingState.getBuiltArchive();
+    if (!newArchive) {
+      this.logger.debug('Archive snapshot not yet captured, skipping header check.');
+      return;
+    }
+
     const header = await buildHeaderFromCircuitOutputs(output);
 
     if (!(await header.hash()).equals(await builtBlockHeader.hash())) {
@@ -480,11 +494,7 @@ export class ProvingOrchestrator implements EpochProver {
       return;
     }
 
-    // Get db for this block
     const blockNumber = provingState.blockNumber;
-    const db = this.dbs.get(blockNumber)!.fork;
-
-    const newArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, db);
     const syncedArchive = await getTreeSnapshot(MerkleTreeId.ARCHIVE, this.dbProvider.getSnapshot(blockNumber));
     if (!syncedArchive.equals(newArchive)) {
       this.logger.error(
@@ -502,12 +512,6 @@ export class ProvingOrchestrator implements EpochProver {
       provingState.reject(`New archive mismatch.`);
       return;
     }
-
-    // TODO(palla/prover): This closes the fork only on the happy path. If this epoch orchestrator
-    // is aborted and never reaches this point, it will leak the fork. We need to add a global cleanup,
-    // but have to make sure it only runs once all operations are completed, otherwise some function here
-    // will attempt to access the fork after it was closed.
-    void this.cleanupDBFork(blockNumber);
   }
 
   /**
@@ -516,6 +520,11 @@ export class ProvingOrchestrator implements EpochProver {
    * If cancelJobsOnStop is false (default), jobs remain in the broker queue and can be reused on restart/reorg.
    */
   public cancel() {
+    void this.deferredJobQueue.cancel();
+    // Recreate the queue so it can accept jobs for subsequent epochs.
+    this.deferredJobQueue = new SerialQueue();
+    this.deferredJobQueue.start(this.enqueueConcurrency);
+
     if (this.cancelJobsOnStop) {
       for (const controller of this.pendingProvingJobs) {
         controller.abort();
@@ -523,6 +532,19 @@ export class ProvingOrchestrator implements EpochProver {
     }
 
     this.provingState?.cancel();
+
+    for (const [blockNumber, db] of this.dbs.entries()) {
+      void db.close().catch(err => this.logger.error(`Error closing db for block ${blockNumber}`, err));
+    }
+    this.dbs.clear();
+  }
+
+  private getDbForBlock(blockNumber: BlockNumber): MerkleTreeWriteOperations {
+    const db = this.dbs.get(blockNumber);
+    if (!db) {
+      throw new Error(`World state fork for block ${blockNumber} not found.`);
+    }
+    return db;
   }
 
   /**
@@ -554,24 +576,6 @@ export class ProvingOrchestrator implements EpochProver {
     return epochProofResult;
   }
 
-  private async cleanupDBFork(blockNumber: BlockNumber): Promise<void> {
-    this.logger.debug(`Cleaning up world state fork for ${blockNumber}`);
-    const fork = this.dbs.get(blockNumber);
-    if (!fork) {
-      return;
-    }
-
-    try {
-      if (!fork.cleanupPromise) {
-        fork.cleanupPromise = fork.fork.close();
-      }
-      await fork.cleanupPromise;
-      this.dbs.delete(blockNumber);
-    } catch (err) {
-      this.logger.error(`Error closing db for block ${blockNumber}`, err);
-    }
-  }
-
   /**
    * Enqueue a job to be scheduled
    * @param provingState - The proving state object being operated on
@@ -630,8 +634,9 @@ export class ProvingOrchestrator implements EpochProver {
       }
     };
 
-    // let the callstack unwind before adding the job to the queue
-    setImmediate(() => void safeJob());
+    // Enqueue onto the serial queue with limited workers to avoid starving the event loop.
+    // Workers yield between jobs via await, allowing I/O callbacks to process.
+    void this.deferredJobQueue.put(() => safeJob());
   }
 
   private async updateL1ToL2MessageTree(l1ToL2Messages: Fr[], db: MerkleTreeWriteOperations) {
@@ -894,17 +899,15 @@ export class ProvingOrchestrator implements EpochProver {
         const leafLocation = provingState.setBlockRootRollupProof(result);
         const checkpointProvingState = provingState.parentCheckpoint;
 
-        // If the proofs were slower than the block header building, then we need to try validating the block header hashes here.
+        // Verification is called from both here and setBlockCompleted. Whichever runs last
+        // will be the first to see all three pieces (header, proof output, archive) and run the checks.
         await this.verifyBuiltBlockAgainstSyncedState(provingState);
 
         if (checkpointProvingState.totalNumBlocks === 1) {
-          this.checkAndEnqueueCheckpointRootRollup(checkpointProvingState);
+          await this.checkAndEnqueueCheckpointRootRollup(checkpointProvingState);
         } else {
-          this.checkAndEnqueueNextBlockMergeRollup(checkpointProvingState, leafLocation);
+          await this.checkAndEnqueueNextBlockMergeRollup(checkpointProvingState, leafLocation);
         }
-
-        // We are finished with the block at this point, ensure the fork is cleaned up
-        void this.cleanupDBFork(provingState.blockNumber);
       },
     );
   }
@@ -1009,14 +1012,14 @@ export class ProvingOrchestrator implements EpochProver {
         },
         signal => this.prover.getBlockMergeRollupProof(inputs, signal, provingState.epochNumber),
       ),
-      result => {
+      async result => {
         provingState.setBlockMergeRollupProof(location, result);
-        this.checkAndEnqueueNextBlockMergeRollup(provingState, location);
+        await this.checkAndEnqueueNextBlockMergeRollup(provingState, location);
       },
     );
   }
 
-  private enqueueCheckpointRootRollup(provingState: CheckpointProvingState) {
+  private async enqueueCheckpointRootRollup(provingState: CheckpointProvingState) {
     if (!provingState.verifyState()) {
       this.logger.debug('Not running checkpoint root rollup. State no longer valid.');
       return;
@@ -1031,7 +1034,7 @@ export class ProvingOrchestrator implements EpochProver {
 
     this.logger.debug(`Enqueuing ${rollupType} for checkpoint ${provingState.index}.`);
 
-    const inputs = provingState.getCheckpointRootRollupInputs();
+    const inputs = await provingState.getCheckpointRootRollupInputs();
 
     this.deferredProving(
       provingState,
@@ -1191,25 +1194,28 @@ export class ProvingOrchestrator implements EpochProver {
     this.enqueueBlockRootRollup(provingState);
   }
 
-  private checkAndEnqueueNextBlockMergeRollup(provingState: CheckpointProvingState, currentLocation: TreeNodeLocation) {
+  private async checkAndEnqueueNextBlockMergeRollup(
+    provingState: CheckpointProvingState,
+    currentLocation: TreeNodeLocation,
+  ) {
     if (!provingState.isReadyForBlockMerge(currentLocation)) {
       return;
     }
 
     const parentLocation = provingState.getParentLocation(currentLocation);
     if (parentLocation.level === 0) {
-      this.checkAndEnqueueCheckpointRootRollup(provingState);
+      await this.checkAndEnqueueCheckpointRootRollup(provingState);
     } else {
       this.enqueueBlockMergeRollup(provingState, parentLocation);
     }
   }
 
-  private checkAndEnqueueCheckpointRootRollup(provingState: CheckpointProvingState) {
+  private async checkAndEnqueueCheckpointRootRollup(provingState: CheckpointProvingState) {
     if (!provingState.isReadyForCheckpointRoot()) {
       return;
     }
 
-    this.enqueueCheckpointRootRollup(provingState);
+    await this.enqueueCheckpointRootRollup(provingState);
   }
 
   private checkAndEnqueueNextCheckpointMergeRollup(provingState: EpochProvingState, currentLocation: TreeNodeLocation) {

package/src/prover-client/prover-client.ts

@@ -54,6 +54,7 @@ export class ProverClient implements EpochProverManager {
       facade,
       this.config.proverId,
       this.config.cancelJobsOnStop,
+      this.config.enqueueConcurrency,
       this.telemetry,
       bindings,
     );
@@ -156,7 +157,7 @@ export class ProverClient implements EpochProverManager {
 }
 
 export function buildServerCircuitProver(
-  config: ActualProverConfig & ACVMConfig & BBConfig,
+  config: Omit<ActualProverConfig, 'enqueueConcurrency'> & ACVMConfig & BBConfig,
   telemetry: TelemetryClient,
 ): Promise<ServerCircuitProver> {
   if (config.realProofs) {

package/src/proving_broker/broker_prover_facade.ts

@@ -5,7 +5,6 @@ import type {
   RECURSIVE_PROOF_LENGTH,
 } from '@aztec/constants';
 import { EpochNumber } from '@aztec/foundation/branded-types';
-import { sha256 } from '@aztec/foundation/crypto/sha256';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { truncate } from '@aztec/foundation/string';
@@ -46,6 +45,8 @@ import type {
   TxRollupPublicInputs,
 } from '@aztec/stdlib/rollup';
 
+import { createHash } from 'node:crypto';
+
 import { InlineProofStore, type ProofStore } from './proof_store/index.js';
 
 // Perform a snapshot sync every 30 seconds
@@ -659,8 +660,12 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
     );
   }
 
-  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }, epochNumber = EpochNumber.ZERO) {
-    const inputsHash = sha256(inputs.toBuffer());
-    return makeProvingJobId(epochNumber, type, inputsHash.toString('hex'));
+  private generateId(
+    type: ProvingRequestType,
+    inputs: { toBuffer(): Buffer },
+    epochNumber = EpochNumber.ZERO,
+  ): ProvingJobId {
+    const inputsHash = createHash('sha256').update(inputs.toBuffer()).digest('hex');
+    return makeProvingJobId(epochNumber, type, inputsHash);
   }
 }

package/src/proving_broker/config.ts

@@ -6,8 +6,8 @@ import {
   numberConfigHelper,
 } from '@aztec/foundation/config';
 import { pickConfigMappings } from '@aztec/foundation/config';
-import { type DataStoreConfig, dataConfigMappings } from '@aztec/kv-store/config';
 import { type ChainConfig, chainConfigMappings } from '@aztec/stdlib/config';
+import { type DataStoreConfig, dataConfigMappings } from '@aztec/stdlib/kv-store';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
 import { z } from 'zod';

package/src/proving_broker/proving_broker.ts

@@ -314,7 +314,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
     // notify listeners of the cancellation
     if (!this.resultsCache.has(id)) {
       this.logger.info(`Cancelling job id=${id}`, { provingJobId: id });
-      await this.#reportProvingJobError(id, 'Aborted', false);
+      await this.#reportProvingJobError(id, 'Aborted', false, undefined, true);
     }
   }
 
@@ -395,6 +395,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
     err: string,
     retry = false,
     filter?: ProvingJobFilter,
+    aborted = false,
   ): Promise<GetProvingJobResponse | undefined> {
     const info = this.inProgress.get(id);
     const item = this.jobsCache.get(id);
@@ -455,7 +456,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
     this.promises.get(id)!.resolve(result);
     this.completedJobNotifications.push(id);
 
-    this.instrumentation.incRejectedJobs(item.type);
+    if (aborted) {
+      this.instrumentation.incAbortedJobs(item.type);
+    } else {
+      this.instrumentation.incRejectedJobs(item.type);
+    }
     if (info) {
       const duration = this.msTimeSource() - info.startedAt;
       this.instrumentation.recordJobDuration(item.type, duration);

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -18,6 +18,7 @@ export class ProvingBrokerInstrumentation {
   private activeJobs: ObservableGauge;
   private resolvedJobs: UpDownCounter;
   private rejectedJobs: UpDownCounter;
+  private abortedJobs: UpDownCounter;
   private timedOutJobs: UpDownCounter;
   private cachedJobs: UpDownCounter;
   private totalJobs: UpDownCounter;
@@ -39,6 +40,8 @@ export class ProvingBrokerInstrumentation {
 
     this.rejectedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_REJECTED_JOBS, provingJobAttrs);
 
+    this.abortedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_ABORTED_JOBS, provingJobAttrs);
+
     this.retriedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RETRIED_JOBS, provingJobAttrs);
 
     this.timedOutJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, provingJobAttrs);
@@ -72,6 +75,12 @@ export class ProvingBrokerInstrumentation {
     });
   }
 
+  incAbortedJobs(proofType: ProvingRequestType) {
+    this.abortedJobs.add(1, {
+      [Attributes.PROVING_JOB_TYPE]: ProvingRequestType[proofType],
+    });
+  }
+
   incRetriedJobs(proofType: ProvingRequestType) {
     this.retriedJobs.add(1, {
       [Attributes.PROVING_JOB_TYPE]: ProvingRequestType[proofType],