@aztec/prover-client 0.69.0-devnet → 0.69.1-devnet

package/src/proving_broker/proving_broker.ts

@@ -9,10 +9,9 @@ import {
   type ProvingJobStatus,
   ProvingRequestType,
 } from '@aztec/circuit-types';
-import { asyncPool } from '@aztec/foundation/async-pool';
 import { createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
-import { PriorityMemoryQueue } from '@aztec/foundation/queue';
+import { PriorityMemoryQueue, SerialQueue } from '@aztec/foundation/queue';
 import { Timer } from '@aztec/foundation/timer';
 import { type TelemetryClient, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
 
@@ -45,7 +44,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   private queues: ProvingQueues = {
     [ProvingRequestType.PUBLIC_VM]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.TUBE_PROOF]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.PRIVATE_KERNEL_EMPTY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
     [ProvingRequestType.PRIVATE_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.PUBLIC_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
@@ -54,6 +52,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
     [ProvingRequestType.BLOCK_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
     [ProvingRequestType.BASE_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
@@ -89,21 +88,24 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   private instrumentation: ProvingBrokerInstrumentation;
   public readonly tracer: Tracer;
 
-  private maxParallelCleanUps: number;
+  private completedJobNotifications: ProvingJobId[] = [];
 
   /**
    * The broker keeps track of the highest epoch its seen.
    * This information is used for garbage collection: once it reaches the next epoch, it can start pruning the database of old state.
-   * This clean up pass is only done against _settled_ jobs. This pass will not cancel jobs that are in-progress or in-queue.
-   * It is a client responsibility to cancel jobs if they are no longer necessary.
+   * It is important that this value is initialised to zero. This ensures that we don't delete any old jobs until the current
+   * process instance receives a job request informing it of the actual current highest epoch
    * Example:
-   * proving epoch 11 - the broker will wipe all setlled jobs for epochs 9 and lower
-   * finished proving epoch 11 and got first job for epoch 12 -> the broker will wipe all setlled jobs for epochs 10 and lower
-   * reorged back to end of epoch 10 -> epoch 11 is skipped and epoch 12 starts -> the broker will wipe all setlled jobs for epochs 10 and lower
+   * proving epoch 11 - the broker will wipe all jobs for epochs 9 and lower
+   * finished proving epoch 11 and got first job for epoch 12 -> the broker will wipe all settled jobs for epochs 10 and lower
+   * reorged back to end of epoch 10 -> epoch 11 is skipped and epoch 12 starts -> the broker will wipe all settled jobs for epochs 10 and lower
    */
   private epochHeight = 0;
   private maxEpochsToKeepResultsFor = 1;
 
+  private requestQueue: SerialQueue = new SerialQueue();
+  private started = false;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     client: TelemetryClient,
@@ -112,7 +114,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       timeoutIntervalMs = 10_000,
       maxRetries = 3,
       maxEpochsToKeepResultsFor = 1,
-      maxParallelCleanUps = 20,
     }: ProofRequestBrokerConfig = {},
     private logger = createLogger('prover-client:proving-broker'),
   ) {
@@ -122,7 +123,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     this.jobTimeoutMs = jobTimeoutMs;
     this.maxRetries = maxRetries;
     this.maxEpochsToKeepResultsFor = maxEpochsToKeepResultsFor;
-    this.maxParallelCleanUps = maxParallelCleanUps;
   }
 
   private measureQueueDepth: MonitorCallback = (type: ProvingRequestType) => {
@@ -142,6 +142,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   };
 
   public start(): Promise<void> {
+    if (this.started) {
+      this.logger.info('Proving Broker already started');
+      return Promise.resolve();
+    }
+    this.logger.info('Proving Broker started');
     for (const [item, result] of this.database.allProvingJobs()) {
       this.logger.info(`Restoring proving job id=${item.id} settled=${!!result}`, {
         provingJobId: item.id,
@@ -161,24 +166,71 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
     this.cleanupPromise.start();
 
+    this.requestQueue.start();
+
     this.instrumentation.monitorQueueDepth(this.measureQueueDepth);
     this.instrumentation.monitorActiveJobs(this.countActiveJobs);
 
+    this.started = true;
+
     return Promise.resolve();
   }
 
   public async stop(): Promise<void> {
+    if (!this.started) {
+      this.logger.warn('ProvingBroker not started');
+      return Promise.resolve();
+    }
+    await this.requestQueue.cancel();
     await this.cleanupPromise.stop();
   }
 
-  public async enqueueProvingJob(job: ProvingJob): Promise<void> {
+  public enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
+    return this.requestQueue.put(() => this.#enqueueProvingJob(job));
+  }
+
+  public cancelProvingJob(id: ProvingJobId): Promise<void> {
+    return this.requestQueue.put(() => this.#cancelProvingJob(id));
+  }
+
+  public getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
+    return this.requestQueue.put(() => this.#getProvingJobStatus(id));
+  }
+
+  public getCompletedJobs(ids: ProvingJobId[]): Promise<ProvingJobId[]> {
+    return this.requestQueue.put(() => this.#getCompletedJobs(ids));
+  }
+
+  public getProvingJob(filter?: ProvingJobFilter): Promise<{ job: ProvingJob; time: number } | undefined> {
+    return this.requestQueue.put(() => this.#getProvingJob(filter));
+  }
+
+  public reportProvingJobSuccess(id: ProvingJobId, value: ProofUri): Promise<void> {
+    return this.requestQueue.put(() => this.#reportProvingJobSuccess(id, value));
+  }
+
+  public reportProvingJobError(id: ProvingJobId, err: string, retry = false): Promise<void> {
+    return this.requestQueue.put(() => this.#reportProvingJobError(id, err, retry));
+  }
+
+  public reportProvingJobProgress(
+    id: ProvingJobId,
+    startedAt: number,
+    filter?: ProvingJobFilter,
+  ): Promise<{ job: ProvingJob; time: number } | undefined> {
+    return this.requestQueue.put(() => this.#reportProvingJobProgress(id, startedAt, filter));
+  }
+
+  async #enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
+    // We return the job status at the start of this call
+    const jobStatus = await this.#getProvingJobStatus(job.id);
     if (this.jobsCache.has(job.id)) {
       const existing = this.jobsCache.get(job.id);
       assert.deepStrictEqual(job, existing, 'Duplicate proving job ID');
       this.logger.debug(`Duplicate proving job id=${job.id} epochNumber=${job.epochNumber}. Ignoring`, {
         provingJobId: job.id,
       });
-      return;
+      return jobStatus;
     }
 
     if (this.isJobStale(job)) {
@@ -199,18 +251,10 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       this.jobsCache.delete(job.id);
       throw err;
     }
+    return jobStatus;
   }
 
-  public waitForJobToSettle(id: ProvingJobId): Promise<ProvingJobSettledResult> {
-    const promiseWithResolvers = this.promises.get(id);
-    if (!promiseWithResolvers) {
-      this.logger.warn(`Job id=${id} not found`, { provingJobId: id });
-      return Promise.resolve({ status: 'rejected', reason: `Job ${id} not found` });
-    }
-    return promiseWithResolvers.promise;
-  }
-
-  public async cancelProvingJob(id: ProvingJobId): Promise<void> {
+  async #cancelProvingJob(id: ProvingJobId): Promise<void> {
     if (!this.jobsCache.has(id)) {
       this.logger.warn(`Can't cancel a job that doesn't exist id=${id}`, { provingJobId: id });
       return;
@@ -219,31 +263,21 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     // 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);
     }
   }
 
-  private async cleanUpProvingJobState(id: ProvingJobId): Promise<void> {
-    if (!this.jobsCache.has(id)) {
-      this.logger.warn(`Can't clean up a job that doesn't exist id=${id}`, { provingJobId: id });
-      return;
-    }
-
-    if (!this.resultsCache.has(id)) {
-      this.logger.warn(`Can't cleanup busy proving job: id=${id}`, { provingJobId: id });
-      return;
+  private cleanUpProvingJobState(ids: ProvingJobId[]) {
+    for (const id of ids) {
+      this.jobsCache.delete(id);
+      this.promises.delete(id);
+      this.resultsCache.delete(id);
+      this.inProgress.delete(id);
+      this.retries.delete(id);
     }
-
-    this.logger.debug(`Cleaning up state for job id=${id}`, { provingJobId: id });
-    await this.database.deleteProvingJobAndResult(id);
-    this.jobsCache.delete(id);
-    this.promises.delete(id);
-    this.resultsCache.delete(id);
-    this.inProgress.delete(id);
-    this.retries.delete(id);
   }
 
-  public getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
+  #getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
     const result = this.resultsCache.get(id);
     if (result) {
       return Promise.resolve(result);
@@ -252,7 +286,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       const item = this.jobsCache.get(id);
 
       if (!item) {
-        this.logger.warn(`Proving job id=${id} not found`, { provingJobId: id });
         return Promise.resolve({ status: 'not-found' });
       }
 
@@ -260,8 +293,15 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     }
   }
 
+  #getCompletedJobs(ids: ProvingJobId[]): Promise<ProvingJobId[]> {
+    const completedJobs = ids.filter(id => this.resultsCache.has(id));
+    const notifications = this.completedJobNotifications;
+    this.completedJobNotifications = [];
+    return Promise.resolve(notifications.concat(completedJobs));
+  }
+
   // eslint-disable-next-line require-await
-  async getProvingJob(
+  async #getProvingJob(
     filter: ProvingJobFilter = { allowList: [] },
   ): Promise<{ job: ProvingJob; time: number } | undefined> {
     const allowedProofs: ProvingRequestType[] =
@@ -299,7 +339,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return undefined;
   }
 
-  async reportProvingJobError(id: ProvingJobId, err: string, retry = false): Promise<void> {
+  async #reportProvingJobError(id: ProvingJobId, err: string, retry = false): Promise<void> {
     const info = this.inProgress.get(id);
     const item = this.jobsCache.get(id);
     const retries = this.retries.get(id) ?? 0;
@@ -351,6 +391,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     const result: ProvingJobSettledResult = { status: 'rejected', reason: String(err) };
     this.resultsCache.set(id, result);
     this.promises.get(id)!.resolve(result);
+    this.completedJobNotifications.push(id);
 
     this.instrumentation.incRejectedJobs(item.type);
     if (info) {
@@ -369,7 +410,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     }
   }
 
-  reportProvingJobProgress(
+  #reportProvingJobProgress(
     id: ProvingJobId,
     startedAt: number,
     filter?: ProvingJobFilter,
@@ -377,12 +418,12 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     const job = this.jobsCache.get(id);
     if (!job) {
       this.logger.warn(`Proving job id=${id} does not exist`, { provingJobId: id });
-      return filter ? this.getProvingJob(filter) : Promise.resolve(undefined);
+      return filter ? this.#getProvingJob(filter) : Promise.resolve(undefined);
     }
 
     if (this.resultsCache.has(id)) {
       this.logger.warn(`Proving job id=${id} has already been completed`, { provingJobId: id });
-      return filter ? this.getProvingJob(filter) : Promise.resolve(undefined);
+      return filter ? this.#getProvingJob(filter) : Promise.resolve(undefined);
     }
 
     const metadata = this.inProgress.get(id);
@@ -420,13 +461,13 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
         } already being worked on by another agent. Sending new one`,
         { provingJobId: id },
       );
-      return this.getProvingJob(filter);
+      return this.#getProvingJob(filter);
     } else {
       return Promise.resolve(undefined);
     }
   }
 
-  async reportProvingJobSuccess(id: ProvingJobId, value: ProofUri): Promise<void> {
+  async #reportProvingJobSuccess(id: ProvingJobId, value: ProofUri): Promise<void> {
     const info = this.inProgress.get(id);
     const item = this.jobsCache.get(id);
     const retries = this.retries.get(id) ?? 0;
@@ -459,6 +500,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     const result: ProvingJobSettledResult = { status: 'fulfilled', value };
     this.resultsCache.set(id, result);
     this.promises.get(id)!.resolve(result);
+    this.completedJobNotifications.push(id);
 
     this.instrumentation.incResolvedJobs(item.type);
     if (info) {
@@ -479,30 +521,32 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
   @trackSpan('ProvingBroker.cleanupPass')
   private async cleanupPass() {
-    await this.cleanupStaleJobs();
-    await this.reEnqueueExpiredJobs();
+    this.cleanupStaleJobs();
+    this.reEnqueueExpiredJobs();
+    const oldestEpochToKeep = this.oldestEpochToKeep();
+    if (oldestEpochToKeep > 0) {
+      await this.requestQueue.put(() => this.database.deleteAllProvingJobsOlderThanEpoch(oldestEpochToKeep));
+      this.logger.trace(`Deleted all epochs older than ${oldestEpochToKeep}`);
+    }
   }
 
-  private async cleanupStaleJobs() {
+  private cleanupStaleJobs() {
     const jobIds = Array.from(this.jobsCache.keys());
     const jobsToClean: ProvingJobId[] = [];
     for (const id of jobIds) {
       const job = this.jobsCache.get(id)!;
-      const isComplete = this.resultsCache.has(id);
-      if (isComplete && this.isJobStale(job)) {
+      if (this.isJobStale(job)) {
         jobsToClean.push(id);
       }
     }
 
     if (jobsToClean.length > 0) {
-      this.logger.info(`Cleaning up jobs=${jobsToClean.length}`);
-      await asyncPool(this.maxParallelCleanUps, jobsToClean, async jobId => {
-        await this.cleanUpProvingJobState(jobId);
-      });
+      this.cleanUpProvingJobState(jobsToClean);
+      this.logger.info(`Cleaned up jobs=${jobsToClean.length}`);
     }
   }
 
-  private async reEnqueueExpiredJobs() {
+  private reEnqueueExpiredJobs() {
     const inProgressEntries = Array.from(this.inProgress.entries());
     for (const [id, metadata] of inProgressEntries) {
       const item = this.jobsCache.get(id);
@@ -515,15 +559,10 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       const now = this.msTimeSource();
       const msSinceLastUpdate = now - metadata.lastUpdatedAt;
       if (msSinceLastUpdate >= this.jobTimeoutMs) {
-        if (this.isJobStale(item)) {
-          // the job has timed out and it's also old, just cancel and move on
-          await this.cancelProvingJob(item.id);
-        } else {
-          this.logger.warn(`Proving job id=${id} timed out. Adding it back to the queue.`, { provingJobId: id });
-          this.inProgress.delete(id);
-          this.enqueueJobInternal(item);
-          this.instrumentation.incTimedOutJobs(item.type);
-        }
+        this.logger.warn(`Proving job id=${id} timed out. Adding it back to the queue.`, { provingJobId: id });
+        this.inProgress.delete(id);
+        this.enqueueJobInternal(item);
+        this.instrumentation.incTimedOutJobs(item.type);
       }
     }
   }
@@ -541,7 +580,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   }
 
   private isJobStale(job: ProvingJob) {
-    return job.epochNumber < this.epochHeight - this.maxEpochsToKeepResultsFor;
+    return job.epochNumber < this.oldestEpochToKeep();
+  }
+
+  private oldestEpochToKeep() {
+    return this.epochHeight - this.maxEpochsToKeepResultsFor;
   }
 }
 
@@ -601,6 +644,7 @@ function proofTypeComparator(a: ProvingRequestType, b: ProvingRequestType): -1 |
  */
 const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
   ProvingRequestType.BLOCK_ROOT_ROLLUP,
+  ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP,
   ProvingRequestType.BLOCK_MERGE_ROLLUP,
   ProvingRequestType.ROOT_ROLLUP,
   ProvingRequestType.MERGE_ROLLUP,
@@ -611,5 +655,4 @@ const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
   ProvingRequestType.ROOT_PARITY,
   ProvingRequestType.BASE_PARITY,
   ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
-  ProvingRequestType.PRIVATE_KERNEL_EMPTY,
 ];

package/src/proving_broker/proving_broker_database/memory.ts

@@ -1,4 +1,10 @@
-import type { ProofUri, ProvingJob, ProvingJobId, ProvingJobSettledResult } from '@aztec/circuit-types';
+import {
+  type ProofUri,
+  type ProvingJob,
+  type ProvingJobId,
+  type ProvingJobSettledResult,
+  getEpochFromProvingJobId,
+} from '@aztec/circuit-types';
 
 import { type ProvingBrokerDatabase } from '../proving_broker_database.js';
 
@@ -29,15 +35,29 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
     return Promise.resolve();
   }
 
-  deleteProvingJobAndResult(id: ProvingJobId): Promise<void> {
-    this.jobs.delete(id);
-    this.results.delete(id);
+  deleteProvingJobs(ids: ProvingJobId[]): Promise<void> {
+    for (const id of ids) {
+      this.jobs.delete(id);
+      this.results.delete(id);
+    }
     return Promise.resolve();
   }
 
+  deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void> {
+    const toDelete = [
+      ...Array.from(this.jobs.keys()).filter(x => getEpochFromProvingJobId(x) < epochNumber),
+      ...Array.from(this.results.keys()).filter(x => getEpochFromProvingJobId(x) < epochNumber),
+    ];
+    return this.deleteProvingJobs(toDelete);
+  }
+
   *allProvingJobs(): Iterable<[ProvingJob, ProvingJobSettledResult | undefined]> {
     for (const item of this.jobs.values()) {
       yield [item, this.results.get(item.id)] as const;
     }
   }
+
+  close(): Promise<void> {
+    return Promise.resolve();
+  }
 }

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -1,27 +1,35 @@
-import { type ProofUri, ProvingJob, type ProvingJobId, ProvingJobSettledResult } from '@aztec/circuit-types';
+import {
+  type ProofUri,
+  ProvingJob,
+  type ProvingJobId,
+  ProvingJobSettledResult,
+  getEpochFromProvingJobId,
+} from '@aztec/circuit-types';
 import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
-import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { type AztecMap } from '@aztec/kv-store';
+import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
 import { Attributes, LmdbMetrics, type TelemetryClient } from '@aztec/telemetry-client';
 
+import { mkdir, readdir } from 'fs/promises';
+import { join } from 'path';
+
+import { type ProverBrokerConfig } from '../config.js';
 import { type ProvingBrokerDatabase } from '../proving_broker_database.js';
 
-export class KVBrokerDatabase implements ProvingBrokerDatabase {
+class SingleEpochDatabase {
   private jobs: AztecMap<ProvingJobId, string>;
   private jobResults: AztecMap<ProvingJobId, string>;
-  private metrics: LmdbMetrics;
 
-  constructor(private store: AztecKVStore, client: TelemetryClient) {
-    this.metrics = new LmdbMetrics(
-      client.getMeter('KVBrokerDatabase'),
-      {
-        [Attributes.DB_DATA_TYPE]: 'prover-broker',
-      },
-      () => store.estimateSize(),
-    );
+  constructor(public readonly store: AztecLmdbStore) {
     this.jobs = store.openMap('proving_jobs');
     this.jobResults = store.openMap('proving_job_results');
   }
 
+  estimateSize() {
+    return this.store.estimateSize();
+  }
+
   async addProvingJob(job: ProvingJob): Promise<void> {
     await this.jobs.set(job.id, jsonStringify(job));
   }
@@ -35,13 +43,6 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     }
   }
 
-  deleteProvingJobAndResult(id: ProvingJobId): Promise<void> {
-    return this.store.transaction(() => {
-      void this.jobs.delete(id);
-      void this.jobResults.delete(id);
-    });
-  }
-
   async setProvingJobError(id: ProvingJobId, reason: string): Promise<void> {
     const result: ProvingJobSettledResult = { status: 'rejected', reason };
     await this.jobResults.set(id, jsonStringify(result));
@@ -51,4 +52,125 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     const result: ProvingJobSettledResult = { status: 'fulfilled', value };
     await this.jobResults.set(id, jsonStringify(result));
   }
+
+  delete() {
+    return this.store.delete();
+  }
+
+  close() {
+    return this.store.close();
+  }
+}
+
+export class KVBrokerDatabase implements ProvingBrokerDatabase {
+  private metrics: LmdbMetrics;
+
+  private constructor(
+    private epochs: Map<number, SingleEpochDatabase>,
+    private config: ProverBrokerConfig,
+    client: TelemetryClient,
+    private logger: Logger,
+  ) {
+    this.metrics = new LmdbMetrics(
+      client.getMeter('KVBrokerDatabase'),
+      {
+        [Attributes.DB_DATA_TYPE]: 'prover-broker',
+      },
+      () => this.estimateSize(),
+    );
+  }
+
+  private estimateSize() {
+    const sizes = Array.from(this.epochs.values()).map(x => x.estimateSize());
+    return {
+      mappingSize: this.config.dataStoreMapSizeKB,
+      numItems: sizes.reduce((prev, curr) => prev + curr.numItems, 0),
+      actualSize: sizes.reduce((prev, curr) => prev + curr.actualSize, 0),
+    };
+  }
+
+  public static async new(
+    config: ProverBrokerConfig,
+    client: TelemetryClient,
+    logger = createLogger('prover-client:proving-broker-database'),
+  ) {
+    const epochs: Map<number, SingleEpochDatabase> = new Map<number, SingleEpochDatabase>();
+    const files = await readdir(config.dataDirectory!, { recursive: false, withFileTypes: true });
+    for (const file of files) {
+      if (!file.isDirectory()) {
+        continue;
+      }
+      const fullDirectory = join(config.dataDirectory!, file.name);
+      const epochDirectory = file.name;
+      const epochNumber = parseInt(epochDirectory, 10);
+      if (!Number.isSafeInteger(epochNumber) || epochNumber < 0) {
+        logger.warn(`Found invalid epoch directory ${fullDirectory} when loading epoch databases, ignoring`);
+        continue;
+      }
+      logger.info(
+        `Loading broker database for epoch ${epochNumber} from ${fullDirectory} with map size ${config.dataStoreMapSizeKB}KB`,
+      );
+      const db = AztecLmdbStore.open(fullDirectory, config.dataStoreMapSizeKB, false);
+      const epochDb = new SingleEpochDatabase(db);
+      epochs.set(epochNumber, epochDb);
+    }
+    return new KVBrokerDatabase(epochs, config, client, logger);
+  }
+
+  async close(): Promise<void> {
+    for (const [_, v] of this.epochs) {
+      await v.close();
+    }
+  }
+
+  async deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void> {
+    const oldEpochs = Array.from(this.epochs.keys()).filter(e => e < epochNumber);
+    for (const old of oldEpochs) {
+      const db = this.epochs.get(old);
+      if (!db) {
+        continue;
+      }
+      this.logger.info(`Deleting broker database for epoch ${old}`);
+      await db.delete();
+      this.epochs.delete(old);
+    }
+  }
+
+  async addProvingJob(job: ProvingJob): Promise<void> {
+    let epochDb = this.epochs.get(job.epochNumber);
+    if (!epochDb) {
+      const newEpochDirectory = join(this.config.dataDirectory!, job.epochNumber.toString());
+      await mkdir(newEpochDirectory, { recursive: true });
+      this.logger.info(
+        `Creating broker database for epoch ${job.epochNumber} at ${newEpochDirectory} with map size ${this.config.dataStoreMapSizeKB}`,
+      );
+      const db = AztecLmdbStore.open(newEpochDirectory, this.config.dataStoreMapSizeKB, false);
+      epochDb = new SingleEpochDatabase(db);
+      this.epochs.set(job.epochNumber, epochDb);
+    }
+    await epochDb.addProvingJob(job);
+  }
+
+  *allProvingJobs(): Iterable<[ProvingJob, ProvingJobSettledResult | undefined]> {
+    const iterators = Array.from(this.epochs.values()).map(x => x.allProvingJobs());
+    for (const it of iterators) {
+      yield* it;
+    }
+  }
+
+  async setProvingJobError(id: ProvingJobId, reason: string): Promise<void> {
+    const epochDb = this.epochs.get(getEpochFromProvingJobId(id));
+    if (!epochDb) {
+      return;
+    }
+    await epochDb.setProvingJobError(id, reason);
+  }
+
+  async setProvingJobResult(id: ProvingJobId, value: ProofUri): Promise<void> {
+    const epochDb = this.epochs.get(getEpochFromProvingJobId(id));
+    if (!epochDb) {
+      return;
+    }
+    await epochDb.setProvingJobResult(id, value);
+  }
 }

package/src/proving_broker/proving_broker_database.ts

@@ -11,10 +11,10 @@ export interface ProvingBrokerDatabase {
   addProvingJob(request: ProvingJob): Promise<void>;
 
   /**
-   * Removes a proof request from the backend
-   * @param id - The ID of the proof request to remove
+   * Deletes all proving jobs belonging to epochs older than the given epoch
+   * @param epochNumber - The epoch number beyond which jobs should be deleted
    */
-  deleteProvingJobAndResult(id: ProvingJobId): Promise<void>;
+  deleteAllProvingJobsOlderThanEpoch(epochNumber: number): Promise<void>;
 
   /**
    * Returns an iterator over all saved proving jobs
@@ -36,4 +36,9 @@ export interface ProvingBrokerDatabase {
    * @param err - The error that occurred while processing the proof request
    */
   setProvingJobError(id: ProvingJobId, err: string): Promise<void>;
+
+  /**
+   * Closes the database
+   */
+  close(): Promise<void>;
 }