@aztec/prover-client 0.0.1-commit.b2a5d0dd1 → 0.0.1-commit.b3d3157a

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

@@ -1,5 +1,7 @@
 import { type ACVMConfig, type BBConfig, BBNativeRollupProver, TestCircuitProver } from '@aztec/bb-prover';
+import type { EpochNumber } from '@aztec/foundation/branded-types';
 import { times } from '@aztec/foundation/collection';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EthAddress } from '@aztec/foundation/eth-address';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { NativeACVMSimulator } from '@aztec/simulator/server';
@@ -15,19 +17,64 @@ import {
   type ServerCircuitProver,
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
+import type { CheckpointConstantData } from '@aztec/stdlib/rollup';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import type { ProverClientConfig } from '../config.js';
+import { CheckpointSubTreeOrchestrator } from '../orchestrator/checkpoint-sub-tree-orchestrator.js';
+import { EpochProvingContext } from '../orchestrator/epoch-proving-context.js';
 import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
+import { TopTreeOrchestrator } from '../orchestrator/top-tree-orchestrator.js';
 import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
 import { InlineProofStore, type ProofStore, createProofStore } from '../proving_broker/proof_store/index.js';
 import { ProvingAgent } from '../proving_broker/proving_agent.js';
 import { ServerEpochProver } from './server-epoch-prover.js';
 
+/**
+ * The factory surface that `EpochProvingJob` (in `prover-node`) depends on. Implemented
+ * by `ProverClient`. Defined here rather than in stdlib because the return types
+ * (`CheckpointSubTreeOrchestrator`, `TopTreeOrchestrator`) are concrete classes from
+ * this package.
+ *
+ * A single `BrokerCircuitProverFacade` is owned by `ProverClient` and shared across
+ * every orchestrator (every sub-tree and every top-tree across every concurrent epoch
+ * job). The broker delivers each completed-job notification exactly once (drained on
+ * the first `getCompletedJobs` poll), so multiple facades polling the same broker
+ * race and lose notifications
+ *
+ * The facade's job map cleans up entries on resolve/reject, and the prover-node
+ * keeps `ProverClient` alive for its whole lifetime
+ */
+export interface EpochProverFactory {
+  getProverId(): EthAddress;
+  /**
+   * Constructs a per-epoch shared context for the caching of e.g. chonk verifier results
+   */
+  createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext;
+  /**
+   * Constructs and starts a `CheckpointSubTreeOrchestrator` for a single checkpoint.
+   */
+  createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator>;
+  createTopTreeOrchestrator(): TopTreeOrchestrator;
+}
+
 /** Manages proving of epochs by orchestrating the proving of individual blocks relying on a pool of prover agents. */
-export class ProverClient implements EpochProverManager {
+export class ProverClient implements EpochProverManager, EpochProverFactory {
   private running = false;
   private agents: ProvingAgent[] = [];
+  /**
+   * The single broker facade shared by every orchestrator created from this client.
+   * Constructed lazily on `start()` and torn down on `stop()` — see the comment on
+   * `EpochProverFactory` for why a single shared facade is required.
+   */
+  private facade: BrokerCircuitProverFacade | undefined;
 
   private constructor(
     private config: ProverClientConfig,
@@ -40,6 +87,38 @@ export class ProverClient implements EpochProverManager {
     private log: Logger = createLogger('prover-client:tx-prover'),
   ) {}
 
+  /**
+   * Lazy-init the shared facade. The broker delivers each completed-job notification
+   * exactly once (drained on the first `getCompletedJobs` poll), so we cannot start
+   * a shared facade alongside the per-call facades that `createEpochProver` builds —
+   * they would race for notifications and one side would silently drop them. Starting
+   * the shared facade only on first use of one of the new factory methods keeps the
+   * legacy `createEpochProver` path race-free.
+   */
+  private getFacade(): BrokerCircuitProverFacade {
+    if (!this.running) {
+      throw new Error('ProverClient is not running; call start() before constructing orchestrators.');
+    }
+    if (!this.facade) {
+      this.facade = new BrokerCircuitProverFacade(
+        this.orchestratorClient,
+        this.proofStore,
+        this.failedProofStore,
+        undefined,
+        this.log.getBindings(),
+      );
+      this.facade.start();
+    }
+    return this.facade;
+  }
+
+  /**
+   * Legacy single-class epoch prover. Each call constructs its own
+   * `BrokerCircuitProverFacade`; the new factory methods (`createCheckpointSubTreeOrchestrator`,
+   * `createTopTreeOrchestrator`, `createEpochProvingContext`) share a single facade
+   * owned by `ProverClient`. Both APIs coexist while the prover-node migrates onto
+   * the new pair.
+   */
   public createEpochProver(): EpochProver {
     const bindings = this.log.getBindings();
     const facade = new BrokerCircuitProverFacade(
@@ -61,6 +140,43 @@ export class ProverClient implements EpochProverManager {
     return new ServerEpochProver(facade, orchestrator);
   }
 
+  public createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext {
+    return new EpochProvingContext(this.getFacade(), epochNumber, this.log.getBindings());
+  }
+
+  public createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator> {
+    return CheckpointSubTreeOrchestrator.start(
+      this.worldState,
+      this.getFacade(),
+      this.config.proverId,
+      epochContext,
+      this.config.cancelJobsOnStop,
+      this.config.enqueueConcurrency,
+      checkpointConstants,
+      l1ToL2Messages,
+      totalNumBlocks,
+      headerOfLastBlockInPreviousCheckpoint,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
+  public createTopTreeOrchestrator(): TopTreeOrchestrator {
+    return new TopTreeOrchestrator(
+      this.getFacade(),
+      this.config.proverId,
+      this.config.enqueueConcurrency,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
   public getProverId(): EthAddress {
     return this.config.proverId;
   }
@@ -100,6 +216,14 @@ export class ProverClient implements EpochProverManager {
     }
     this.running = false;
     await this.stopAgents();
+    if (this.facade) {
+      try {
+        await this.facade.stop();
+      } catch (err) {
+        this.log.error('Error stopping shared broker facade', err);
+      }
+      this.facade = undefined;
+    }
     await tryStop(this.orchestratorClient);
   }
 

package/src/proving_broker/config.ts

@@ -4,6 +4,7 @@ import {
   booleanConfigHelper,
   getDefaultConfig,
   numberConfigHelper,
+  optionalNumberConfigHelper,
 } from '@aztec/foundation/config';
 import { pickConfigMappings } from '@aztec/foundation/config';
 import { type ChainConfig, chainConfigMappings } from '@aztec/stdlib/config';
@@ -73,7 +74,7 @@ export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig>
   },
   proverBrokerStoreMapSizeKb: {
     env: 'PROVER_BROKER_STORE_MAP_SIZE_KB',
-    parseEnv: (val: string) => +val,
+    ...optionalNumberConfigHelper(),
     description: "The size of the prover broker's database. Will override the dataStoreMapSizeKb if set.",
   },
   proverBrokerDebugReplayEnabled: {

package/src/proving_broker/index.ts

@@ -7,3 +7,4 @@ export * from './proving_broker_database/persisted.js';
 export * from './proof_store/index.js';
 export * from './factory.js';
 export * from './config.js';
+export { BrokerCircuitProverFacade } from './broker_prover_facade.js';

package/src/proving_broker/proving_broker.ts

@@ -319,6 +319,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
   }
 
   private cleanUpProvingJobState(ids: ProvingJobId[]) {
+    const idsToClean = new Set(ids);
     for (const id of ids) {
       this.jobsCache.delete(id);
       const deferred = this.promises.get(id);
@@ -331,6 +332,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
       this.retries.delete(id);
       this.enqueuedAt.delete(id);
     }
+    this.completedJobNotifications = this.completedJobNotifications.filter(id => !idsToClean.has(id));
   }
 
   #getProvingJobStatus(id: ProvingJobId): ProvingJobStatus {
@@ -598,21 +600,21 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
   }
 
   private async cleanupPass() {
-    this.cleanupStaleJobs();
     this.reEnqueueExpiredJobs();
     const oldestEpochToKeep = this.oldestEpochToKeep();
     if (oldestEpochToKeep > 0) {
+      this.cleanupJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       await this.database.deleteAllProvingJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       this.logger.trace(`Deleted all epochs older than ${oldestEpochToKeep}`);
     }
   }
 
-  private cleanupStaleJobs() {
+  private cleanupJobsOlderThanEpoch(epochNumber: EpochNumber) {
     const jobIds = Array.from(this.jobsCache.keys());
     const jobsToClean: ProvingJobId[] = [];
     for (const id of jobIds) {
       const job = this.jobsCache.get(id)!;
-      if (this.isJobStale(job) && !this.inProgress.has(id) && !this.resultsCache.has(id)) {
+      if (job.epochNumber < epochNumber) {
         jobsToClean.push(id);
       }
     }

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -154,7 +154,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         fullDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        config.l1Contracts.rollupAddress,
+        config.rollupAddress,
         config.dataStoreMapSizeKb,
       );
       const epochDb = new SingleEpochDatabase(db);
@@ -222,7 +222,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         newEpochDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        this.config.l1Contracts.rollupAddress,
+        this.config.rollupAddress,
         this.config.dataStoreMapSizeKb,
       );
       epochDb = new SingleEpochDatabase(db);

package/src/proving_broker/rpc.ts

@@ -37,26 +37,29 @@ const GetProvingJobResponse = z.object({
 });
 
 export const ProvingJobProducerSchema: ApiSchemaFor<ProvingJobProducer> = {
-  enqueueProvingJob: z.function().args(ProvingJob).returns(ProvingJobStatus),
-  getProvingJobStatus: z.function().args(ProvingJobId).returns(ProvingJobStatus),
-  cancelProvingJob: z.function().args(ProvingJobId).returns(z.void()),
-  getCompletedJobs: z.function().args(z.array(ProvingJobId)).returns(z.array(ProvingJobId)),
+  enqueueProvingJob: z.function({ input: z.tuple([ProvingJob]), output: ProvingJobStatus }),
+  getProvingJobStatus: z.function({ input: z.tuple([ProvingJobId]), output: ProvingJobStatus }),
+  cancelProvingJob: z.function({ input: z.tuple([ProvingJobId]), output: z.void() }),
+  getCompletedJobs: z.function({ input: z.tuple([z.array(ProvingJobId)]), output: z.array(ProvingJobId) }),
 };
 
 export const ProvingJobConsumerSchema: ApiSchemaFor<ProvingJobConsumer> = {
-  getProvingJob: z.function().args(optional(ProvingJobFilterSchema)).returns(GetProvingJobResponse.optional()),
-  reportProvingJobError: z
-    .function()
-    .args(ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobProgress: z
-    .function()
-    .args(ProvingJobId, z.number(), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobSuccess: z
-    .function()
-    .args(ProvingJobId, ProofUri, optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
+  getProvingJob: z.function({
+    input: z.tuple([optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobError: z.function({
+    input: z.tuple([ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobProgress: z.function({
+    input: z.tuple([ProvingJobId, z.number(), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobSuccess: z.function({
+    input: z.tuple([ProvingJobId, ProofUri, optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
 };
 
 export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
@@ -65,10 +68,10 @@ export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
 };
 
 export const ProvingJobBrokerDebugSchema: ApiSchemaFor<ProvingJobBrokerDebug> = {
-  replayProvingJob: z
-    .function()
-    .args(ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri)
-    .returns(ProvingJobStatus),
+  replayProvingJob: z.function({
+    input: z.tuple([ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri]),
+    output: ProvingJobStatus,
+  }),
 };
 
 export const ProvingJobBrokerSchemaWithDebug: ApiSchemaFor<ProvingJobBroker & ProvingJobBrokerDebug> = {