@aztec/prover-client 0.66.0 → 0.67.0

package/src/proving_broker/caching_broker_facade.ts

@@ -33,7 +33,7 @@ import {
   type TubeInputs,
 } from '@aztec/circuits.js';
 import { sha256 } from '@aztec/foundation/crypto';
-import { createDebugLogger } from '@aztec/foundation/log';
+import { createLogger } from '@aztec/foundation/log';
 import { retryUntil } from '@aztec/foundation/retry';
 
 import { InlineProofStore, type ProofStore } from './proof_store.js';
@@ -52,13 +52,14 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     private proofStore: ProofStore = new InlineProofStore(),
     private waitTimeoutMs = MAX_WAIT_MS,
     private pollIntervalMs = 1000,
-    private log = createDebugLogger('aztec:prover-client:caching-prover-broker'),
+    private log = createLogger('prover-client:caching-prover-broker'),
   ) {}
 
   private async enqueueAndWaitForJob<T extends ProvingRequestType>(
     id: ProvingJobId,
     type: T,
     inputs: ProvingJobInputsMap[T],
+    epochNumber = 0,
     signal?: AbortSignal,
   ): Promise<ProvingJobResultsMap[T]> {
     // first try the cache
@@ -95,6 +96,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
           id,
           type,
           inputsUri,
+          epochNumber,
         });
         await this.cache.setProvingJobStatus(id, { status: 'in-queue' });
       } catch (err) {
@@ -107,7 +109,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     // notify broker of cancelled job
     const abortFn = async () => {
       signal?.removeEventListener('abort', abortFn);
-      await this.broker.removeAndCancelProvingJob(id);
+      await this.broker.cancelProvingJob(id);
     };
 
     signal?.addEventListener('abort', abortFn);
@@ -147,18 +149,21 @@ export class CachingBrokerFacade implements ServerCircuitProver {
       }
     } finally {
       signal?.removeEventListener('abort', abortFn);
+      // we've saved the result in our cache. We can tell the broker to clear its state
+      await this.broker.cleanUpProvingJobState(id);
     }
   }
 
   getAvmProof(
     inputs: AvmCircuitInputs,
     signal?: AbortSignal,
-    _blockNumber?: number,
+    epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof AVM_PROOF_LENGTH_IN_FIELDS>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.PUBLIC_VM, inputs),
       ProvingRequestType.PUBLIC_VM,
       inputs,
+      epochNumber,
       signal,
     );
   }
@@ -166,12 +171,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getBaseParityProof(
     inputs: BaseParityInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.BASE_PARITY, inputs),
       ProvingRequestType.BASE_PARITY,
       inputs,
+      epochNumber,
       signal,
     );
   }
@@ -179,12 +185,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getBlockMergeRollupProof(
     input: BlockMergeRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input),
       ProvingRequestType.BLOCK_MERGE_ROLLUP,
       input,
+      epochNumber,
       signal,
     );
   }
@@ -192,12 +199,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getBlockRootRollupProof(
     input: BlockRootRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input),
       ProvingRequestType.BLOCK_ROOT_ROLLUP,
       input,
+      epochNumber,
       signal,
     );
   }
@@ -205,12 +213,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getEmptyBlockRootRollupProof(
     input: EmptyBlockRootRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input),
       ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
       input,
+      epochNumber,
       signal,
     );
   }
@@ -218,12 +227,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getEmptyPrivateKernelProof(
     inputs: PrivateKernelEmptyInputData,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs),
       ProvingRequestType.PRIVATE_KERNEL_EMPTY,
       inputs,
+      epochNumber,
       signal,
     );
   }
@@ -231,24 +241,26 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getMergeRollupProof(
     input: MergeRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.MERGE_ROLLUP, input),
       ProvingRequestType.MERGE_ROLLUP,
       input,
+      epochNumber,
       signal,
     );
   }
   getPrivateBaseRollupProof(
     baseRollupInput: PrivateBaseRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput),
       ProvingRequestType.PRIVATE_BASE_ROLLUP,
       baseRollupInput,
+      epochNumber,
       signal,
     );
   }
@@ -256,12 +268,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getPublicBaseRollupProof(
     inputs: PublicBaseRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs),
       ProvingRequestType.PUBLIC_BASE_ROLLUP,
       inputs,
+      epochNumber,
       signal,
     );
   }
@@ -269,12 +282,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getRootParityProof(
     inputs: RootParityInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.ROOT_PARITY, inputs),
       ProvingRequestType.ROOT_PARITY,
       inputs,
+      epochNumber,
       signal,
     );
   }
@@ -282,12 +296,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getRootRollupProof(
     input: RootRollupInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<RootRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.ROOT_ROLLUP, input),
       ProvingRequestType.ROOT_ROLLUP,
       input,
+      epochNumber,
       signal,
     );
   }
@@ -295,12 +310,13 @@ export class CachingBrokerFacade implements ServerCircuitProver {
   getTubeProof(
     tubeInput: TubeInputs,
     signal?: AbortSignal,
-    _epochNumber?: number,
+    epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof TUBE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
       this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput),
       ProvingRequestType.TUBE_PROOF,
       tubeInput,
+      epochNumber,
       signal,
     );
   }

package/src/proving_broker/proving_agent.ts

@@ -8,7 +8,7 @@ import {
   ProvingRequestType,
   type ServerCircuitProver,
 } from '@aztec/circuit-types';
-import { createDebugLogger } from '@aztec/foundation/log';
+import { createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { Timer } from '@aztec/foundation/timer';
 import { type TelemetryClient } from '@aztec/telemetry-client';
@@ -39,7 +39,7 @@ export class ProvingAgent {
     private proofAllowList: Array<ProvingRequestType> = [],
     /** How long to wait between jobs */
     private pollIntervalMs = 1000,
-    private log = createDebugLogger('aztec:prover-client:proving-agent'),
+    private log = createLogger('prover-client:proving-agent'),
   ) {
     this.instrumentation = new ProvingAgentInstrumentation(client);
     this.runningPromise = new RunningPromise(this.safeWork, this.pollIntervalMs);

package/src/proving_broker/proving_broker.ts

@@ -9,7 +9,8 @@ import {
   type ProvingJobStatus,
   ProvingRequestType,
 } from '@aztec/circuit-types';
-import { createDebugLogger } from '@aztec/foundation/log';
+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 { Timer } from '@aztec/foundation/timer';
@@ -30,29 +31,33 @@ type ProofRequestBrokerConfig = {
   timeoutIntervalMs?: number;
   jobTimeoutMs?: number;
   maxRetries?: number;
+  maxEpochsToKeepResultsFor?: number;
+  maxParallelCleanUps?: number;
 };
 
+type EnqueuedProvingJob = Pick<ProvingJob, 'id' | 'epochNumber'>;
+
 /**
  * A broker that manages proof requests and distributes them to workers based on their priority.
  * It takes a backend that is responsible for storing and retrieving proof requests and results.
  */
 export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
   private queues: ProvingQueues = {
-    [ProvingRequestType.PUBLIC_VM]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.TUBE_PROOF]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.PRIVATE_KERNEL_EMPTY]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
+    [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<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.PUBLIC_BASE_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.MERGE_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.ROOT_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
+    [ProvingRequestType.PRIVATE_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PUBLIC_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.BLOCK_MERGE_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.BASE_PARITY]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
-    [ProvingRequestType.ROOT_PARITY]: new PriorityMemoryQueue<ProvingJob>(provingJobComparator),
+    [ProvingRequestType.BASE_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.ROOT_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
   };
 
   // holds a copy of the database in memory in order to quickly fulfill requests
@@ -76,23 +81,46 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
   // a map of promises that will be resolved when a job is settled
   private promises = new Map<ProvingJobId, PromiseWithResolvers<ProvingJobSettledResult>>();
 
-  private timeoutPromise: RunningPromise;
-  private timeSource = () => Math.floor(Date.now() / 1000);
+  private cleanupPromise: RunningPromise;
+  private msTimeSource = () => Date.now();
   private jobTimeoutMs: number;
   private maxRetries: number;
 
   private instrumentation: ProvingBrokerInstrumentation;
 
+  private maxParallelCleanUps: number;
+
+  /**
+   * 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.
+   * 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
+   */
+  private epochHeight = 0;
+  private maxEpochsToKeepResultsFor = 1;
+
   public constructor(
     private database: ProvingBrokerDatabase,
     client: TelemetryClient,
-    { jobTimeoutMs = 30_000, timeoutIntervalMs = 10_000, maxRetries = 3 }: ProofRequestBrokerConfig = {},
-    private logger = createDebugLogger('aztec:prover-client:proving-broker'),
+    {
+      jobTimeoutMs = 30_000,
+      timeoutIntervalMs = 10_000,
+      maxRetries = 3,
+      maxEpochsToKeepResultsFor = 1,
+      maxParallelCleanUps = 20,
+    }: ProofRequestBrokerConfig = {},
+    private logger = createLogger('prover-client:proving-broker'),
   ) {
     this.instrumentation = new ProvingBrokerInstrumentation(client);
-    this.timeoutPromise = new RunningPromise(this.timeoutCheck, timeoutIntervalMs);
+    this.cleanupPromise = new RunningPromise(this.cleanupPass, timeoutIntervalMs);
     this.jobTimeoutMs = jobTimeoutMs;
     this.maxRetries = maxRetries;
+    this.maxEpochsToKeepResultsFor = maxEpochsToKeepResultsFor;
+    this.maxParallelCleanUps = maxParallelCleanUps;
   }
 
   private measureQueueDepth: MonitorCallback = (type: ProvingRequestType) => {
@@ -127,7 +155,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       }
     }
 
-    this.timeoutPromise.start();
+    this.cleanupPromise.start();
 
     this.instrumentation.monitorQueueDepth(this.measureQueueDepth);
     this.instrumentation.monitorActiveJobs(this.countActiveJobs);
@@ -135,8 +163,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
     return Promise.resolve();
   }
 
-  public stop(): Promise<void> {
-    return this.timeoutPromise.stop();
+  public async stop(): Promise<void> {
+    await this.cleanupPromise.stop();
   }
 
   public async enqueueProvingJob(job: ProvingJob): Promise<void> {
@@ -159,15 +187,22 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
     return promiseWithResolvers.promise;
   }
 
-  public async removeAndCancelProvingJob(id: ProvingJobId): Promise<void> {
-    this.logger.info(`Cancelling job id=${id}`);
-    await this.database.deleteProvingJobAndResult(id);
-
+  public async cancelProvingJob(id: ProvingJobId): Promise<void> {
     // notify listeners of the cancellation
     if (!this.resultsCache.has(id)) {
-      this.promises.get(id)?.resolve({ status: 'rejected', reason: 'Aborted' });
+      this.logger.info(`Cancelling job id=${id}`);
+      await this.reportProvingJobError(id, 'Aborted', false);
+    }
+  }
+
+  public async cleanUpProvingJobState(id: ProvingJobId): Promise<void> {
+    if (!this.resultsCache.has(id)) {
+      this.logger.warn(`Can't cleanup busy proving job: id=${id}`);
+      return;
     }
 
+    this.logger.debug(`Cleaning up state for job id=${id}`);
+    await this.database.deleteProvingJobAndResult(id);
     this.jobsCache.delete(id);
     this.promises.delete(id);
     this.resultsCache.delete(id);
@@ -204,14 +239,15 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
 
     for (const proofType of allowedProofs) {
       const queue = this.queues[proofType];
-      let job: ProvingJob | undefined;
+      let enqueuedJob: EnqueuedProvingJob | undefined;
       // exhaust the queue and make sure we're not sending a job that's already in progress
       // or has already been completed
       // this can happen if the broker crashes and restarts
       // it's possible agents will report progress or results for jobs that are in the queue (after the restart)
-      while ((job = queue.getImmediate())) {
-        if (!this.inProgress.has(job.id) && !this.resultsCache.has(job.id)) {
-          const time = this.timeSource();
+      while ((enqueuedJob = queue.getImmediate())) {
+        const job = this.jobsCache.get(enqueuedJob.id);
+        if (job && !this.inProgress.has(enqueuedJob.id) && !this.resultsCache.has(enqueuedJob.id)) {
+          const time = this.msTimeSource();
           this.inProgress.set(job.id, {
             id: job.id,
             startedAt: time,
@@ -246,7 +282,12 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       this.inProgress.delete(id);
     }
 
-    if (retry && retries + 1 < this.maxRetries) {
+    if (this.resultsCache.has(id)) {
+      this.logger.warn(`Proving job id=${id} already is already settled, ignoring error`);
+      return;
+    }
+
+    if (retry && retries + 1 < this.maxRetries && !this.isJobStale(item)) {
       this.logger.info(`Retrying proving job id=${id} type=${ProvingRequestType[item.type]} retry=${retries + 1}`);
       this.retries.set(id, retries + 1);
       this.enqueueJobInternal(item);
@@ -254,8 +295,10 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       return;
     }
 
-    this.logger.debug(
-      `Marking proving job id=${id} type=${ProvingRequestType[item.type]} totalAttempts=${retries + 1} as failed`,
+    this.logger.warn(
+      `Marking proving job as failed id=${id} type=${ProvingRequestType[item.type]} totalAttempts=${
+        retries + 1
+      } err=${err}`,
     );
 
     await this.database.setProvingJobError(id, err);
@@ -265,8 +308,8 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
     this.promises.get(id)!.resolve(result);
     this.instrumentation.incRejectedJobs(item.type);
     if (info) {
-      const duration = this.timeSource() - info.startedAt;
-      this.instrumentation.recordJobDuration(item.type, duration * 1000);
+      const duration = this.msTimeSource() - info.startedAt;
+      this.instrumentation.recordJobDuration(item.type, duration);
     }
   }
 
@@ -281,8 +324,13 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       return filter ? this.getProvingJob(filter) : Promise.resolve(undefined);
     }
 
+    if (this.resultsCache.has(id)) {
+      this.logger.warn(`Proving job id=${id} has already been completed`);
+      return filter ? this.getProvingJob(filter) : Promise.resolve(undefined);
+    }
+
     const metadata = this.inProgress.get(id);
-    const now = this.timeSource();
+    const now = this.msTimeSource();
     if (!metadata) {
       this.logger.warn(
         `Proving job id=${id} type=${ProvingRequestType[job.type]} not found in the in-progress cache, adding it`,
@@ -293,7 +341,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       this.inProgress.set(id, {
         id,
         startedAt,
-        lastUpdatedAt: this.timeSource(),
+        lastUpdatedAt: this.msTimeSource(),
       });
       return Promise.resolve(undefined);
     } else if (startedAt <= metadata.startedAt) {
@@ -334,6 +382,11 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
       this.inProgress.delete(id);
     }
 
+    if (this.resultsCache.has(id)) {
+      this.logger.warn(`Proving job id=${id} already settled, ignoring result`);
+      return;
+    }
+
     this.logger.debug(
       `Proving job complete id=${id} type=${ProvingRequestType[item.type]} totalAttempts=${retries + 1}`,
     );
@@ -346,7 +399,31 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
     this.instrumentation.incResolvedJobs(item.type);
   }
 
-  private timeoutCheck = () => {
+  private cleanupPass = async () => {
+    await this.cleanupStaleJobs();
+    await this.reEnqueueExpiredJobs();
+  };
+
+  private async 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)) {
+        jobsToClean.push(id);
+      }
+    }
+
+    if (jobsToClean.length > 0) {
+      this.logger.info(`Cleaning up [${jobsToClean.join(',')}]`);
+      await asyncPool(this.maxParallelCleanUps, jobsToClean, async jobId => {
+        await this.cleanUpProvingJobState(jobId);
+      });
+    }
+  }
+
+  private async reEnqueueExpiredJobs() {
     const inProgressEntries = Array.from(this.inProgress.entries());
     for (const [id, metadata] of inProgressEntries) {
       const item = this.jobsCache.get(id);
@@ -356,28 +433,42 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer {
         continue;
       }
 
-      const msSinceLastUpdate = (this.timeSource() - metadata.lastUpdatedAt) * 1000;
+      const now = this.msTimeSource();
+      const msSinceLastUpdate = now - metadata.lastUpdatedAt;
       if (msSinceLastUpdate >= this.jobTimeoutMs) {
-        this.logger.warn(`Proving job id=${id} timed out. Adding it back to the queue.`);
-        this.inProgress.delete(id);
-        this.enqueueJobInternal(item);
-        this.instrumentation.incTimedOutJobs(item.type);
+        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.`);
+          this.inProgress.delete(id);
+          this.enqueueJobInternal(item);
+          this.instrumentation.incTimedOutJobs(item.type);
+        }
       }
     }
-  };
+  }
 
   private enqueueJobInternal(job: ProvingJob): void {
     if (!this.promises.has(job.id)) {
       this.promises.set(job.id, promiseWithResolvers());
     }
-    this.queues[job.type].put(job);
+    this.queues[job.type].put({
+      epochNumber: job.epochNumber,
+      id: job.id,
+    });
     this.enqueuedAt.set(job.id, new Timer());
+    this.epochHeight = Math.max(this.epochHeight, job.epochNumber);
     this.logger.debug(`Enqueued new proving job id=${job.id}`);
   }
+
+  private isJobStale(job: ProvingJob) {
+    return job.epochNumber < this.epochHeight - this.maxEpochsToKeepResultsFor;
+  }
 }
 
 type ProvingQueues = {
-  [K in ProvingRequestType]: PriorityMemoryQueue<ProvingJob>;
+  [K in ProvingRequestType]: PriorityMemoryQueue<EnqueuedProvingJob>;
 };
 
 /**
@@ -386,12 +477,10 @@ type ProvingQueues = {
  * @param b - Another proving job
  * @returns A number indicating the relative priority of the two proving jobs
  */
-function provingJobComparator(a: ProvingJob, b: ProvingJob): -1 | 0 | 1 {
-  const aBlockNumber = a.blockNumber ?? 0;
-  const bBlockNumber = b.blockNumber ?? 0;
-  if (aBlockNumber < bBlockNumber) {
+function provingJobComparator(a: EnqueuedProvingJob, b: EnqueuedProvingJob): -1 | 0 | 1 {
+  if (a.epochNumber < b.epochNumber) {
     return -1;
-  } else if (aBlockNumber > bBlockNumber) {
+  } else if (a.epochNumber > b.epochNumber) {
     return 1;
   } else {
     return 0;

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -1,7 +1,7 @@
 import { type ProofUri, ProvingJob, type ProvingJobId, ProvingJobSettledResult } from '@aztec/circuit-types';
 import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
 import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
-import { LmdbMetrics, Metrics, type TelemetryClient } from '@aztec/telemetry-client';
+import { Attributes, LmdbMetrics, type TelemetryClient } from '@aztec/telemetry-client';
 
 import { type ProvingBrokerDatabase } from '../proving_broker_database.js';
 
@@ -14,14 +14,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
     this.metrics = new LmdbMetrics(
       client.getMeter('KVBrokerDatabase'),
       {
-        name: Metrics.PROVING_QUEUE_DB_MAP_SIZE,
-        description: 'Database map size for the proving broker',
+        [Attributes.DB_DATA_TYPE]: 'prover-broker',
       },
-      {
-        name: Metrics.PROVING_QUEUE_DB_USED_SIZE,
-        description: 'Database used size for the proving broker',
-      },
-      { name: Metrics.PROVING_QUEUE_DB_NUM_ITEMS, description: 'Number of items in the broker database' },
       () => store.estimateSize(),
     );
     this.jobs = store.openMap('proving_jobs');

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -9,7 +9,6 @@ import {
   type TelemetryClient,
   type UpDownCounter,
   ValueType,
-  millisecondBuckets,
 } from '@aztec/telemetry-client';
 
 export type MonitorCallback = (proofType: ProvingRequestType) => number;
@@ -55,18 +54,12 @@ export class ProvingBrokerInstrumentation {
       description: 'Records how long a job sits in the queue',
       unit: 'ms',
       valueType: ValueType.INT,
-      advice: {
-        explicitBucketBoundaries: millisecondBuckets(1), // 10ms -> ~327s
-      },
     });
 
     this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION, {
       description: 'Records how long a job takes to complete',
       unit: 'ms',
       valueType: ValueType.INT,
-      advice: {
-        explicitBucketBoundaries: millisecondBuckets(1), // 10ms -> ~327s
-      },
     });
   }
 

package/src/proving_broker/rpc.ts

@@ -28,7 +28,8 @@ const GetProvingJobResponse = z.object({
 export const ProvingJobProducerSchema: ApiSchemaFor<ProvingJobProducer> = {
   enqueueProvingJob: z.function().args(ProvingJob).returns(z.void()),
   getProvingJobStatus: z.function().args(ProvingJobId).returns(ProvingJobStatus),
-  removeAndCancelProvingJob: z.function().args(ProvingJobId).returns(z.void()),
+  cleanUpProvingJobState: z.function().args(ProvingJobId).returns(z.void()),
+  cancelProvingJob: z.function().args(ProvingJobId).returns(z.void()),
   waitForJobToSettle: z.function().args(ProvingJobId).returns(ProvingJobSettledResult),
 };
 

package/src/test/mock_prover.ts

@@ -82,8 +82,11 @@ export class TestBroker implements ProvingJobProducer {
   getProvingJobStatus(id: ProvingJobId): Promise<ProvingJobStatus> {
     return this.broker.getProvingJobStatus(id);
   }
-  removeAndCancelProvingJob(id: ProvingJobId): Promise<void> {
-    return this.broker.removeAndCancelProvingJob(id);
+  cleanUpProvingJobState(id: ProvingJobId): Promise<void> {
+    return this.broker.cleanUpProvingJobState(id);
+  }
+  cancelProvingJob(id: string): Promise<void> {
+    return this.broker.cancelProvingJob(id);
   }
   waitForJobToSettle(id: ProvingJobId): Promise<ProvingJobSettledResult> {
     return this.broker.waitForJobToSettle(id);