fix: memory leak in the broker

yarn-project/circuit-types/src/interfaces/prover-broker.ts

@@ -56,10 +56,16 @@ export interface ProvingJobProducer {
   enqueueProvingJob(job: ProvingJob): Promise<void>;
 
   /**
-   * Cancels a proving job and clears all of its
+   * Cancels a proving job.
    * @param id - The ID of the job to cancel
    */
-  removeAndCancelProvingJob(id: ProvingJobId): Promise<void>;
+  cancelProvingJob(id: ProvingJobId): Promise<void>;
+
+  /**
+   * Cleans up after a job has completed. Throws if the job is in-progress
+   * @param id - The ID of the job to cancel
+   */
+  cleanUpProvingJobState(id: ProvingJobId): Promise<void>;
 
   /**
    * Returns the current status fof the proving job

yarn-project/circuit-types/src/interfaces/proving-job-source.test.ts

@@ -70,6 +70,7 @@ class MockProvingJobSource implements ProvingJobSource {
       id: 'a-job-id',
       type: ProvingRequestType.PRIVATE_BASE_ROLLUP,
       inputsUri: 'inputs-uri' as ProofUri,
+      epochNumber: 1,
     });
   }
   heartbeat(jobId: string): Promise<void> {

yarn-project/circuit-types/src/interfaces/proving-job.ts

@@ -241,7 +241,7 @@ export type ProvingJobId = z.infer<typeof ProvingJobId>;
 export const ProvingJob = z.object({
   id: ProvingJobId,
   type: z.nativeEnum(ProvingRequestType),
-  blockNumber: z.number().optional(),
+  epochNumber: z.number(),
   inputsUri: ProofUri,
 });
 

yarn-project/prover-client/src/prover-agent/memory-proving-queue.ts

@@ -120,6 +120,7 @@ export class MemoryProvingQueue implements ServerCircuitProver, ProvingJobSource
         id: job.id,
         type: job.type,
         inputsUri: job.inputsUri,
+        epochNumber: job.epochNumber,
       };
     } catch (err) {
       if (err instanceof TimeoutError) {
@@ -244,7 +245,7 @@ export class MemoryProvingQueue implements ServerCircuitProver, ProvingJobSource
       reject,
       attempts: 1,
       heartbeat: 0,
-      epochNumber,
+      epochNumber: epochNumber ?? 0,
     };
 
     if (signal) {

yarn-project/prover-client/src/proving_broker/caching_broker_facade.test.ts

@@ -24,7 +24,8 @@ describe('CachingBrokerFacade', () => {
     broker = mock<ProvingJobProducer>({
       enqueueProvingJob: jest.fn<any>(),
       getProvingJobStatus: jest.fn<any>(),
-      removeAndCancelProvingJob: jest.fn<any>(),
+      cancelProvingJob: jest.fn<any>(),
+      cleanUpProvingJobState: jest.fn<any>(),
       waitForJobToSettle: jest.fn<any>(),
     });
     cache = new InMemoryProverCache();
@@ -101,4 +102,55 @@ describe('CachingBrokerFacade', () => {
     await expect(facade.getBaseParityProof(inputs)).resolves.toEqual(result);
     expect(broker.enqueueProvingJob).toHaveBeenCalledTimes(1); // job was only ever enqueued once
   });
+
+  it('clears broker state after a job resolves', async () => {
+    const { promise, resolve } = promiseWithResolvers<any>();
+    broker.enqueueProvingJob.mockResolvedValue(Promise.resolve());
+    broker.waitForJobToSettle.mockResolvedValue(promise);
+
+    const inputs = makeBaseParityInputs();
+    void facade.getBaseParityProof(inputs);
+    await jest.advanceTimersToNextTimerAsync();
+
+    const job = broker.enqueueProvingJob.mock.calls[0][0];
+    const result = makePublicInputsAndRecursiveProof(
+      makeParityPublicInputs(),
+      makeRecursiveProof(RECURSIVE_PROOF_LENGTH),
+      VerificationKeyData.makeFakeHonk(),
+    );
+    const outputUri = await proofStore.saveProofOutput(job.id, ProvingRequestType.BASE_PARITY, result);
+    resolve({
+      status: 'fulfilled',
+      value: outputUri,
+    });
+
+    await jest.advanceTimersToNextTimerAsync();
+    expect(broker.cleanUpProvingJobState).toHaveBeenCalled();
+  });
+
+  it('clears broker state after a job is canceled', async () => {
+    const { promise, resolve } = promiseWithResolvers<any>();
+    const catchSpy = jest.fn();
+    broker.enqueueProvingJob.mockResolvedValue(Promise.resolve());
+    broker.waitForJobToSettle.mockResolvedValue(promise);
+
+    const inputs = makeBaseParityInputs();
+    const controller = new AbortController();
+    void facade.getBaseParityProof(inputs, controller.signal).catch(catchSpy);
+    await jest.advanceTimersToNextTimerAsync();
+
+    expect(broker.cancelProvingJob).not.toHaveBeenCalled();
+    controller.abort();
+    await jest.advanceTimersToNextTimerAsync();
+    expect(broker.cancelProvingJob).toHaveBeenCalled();
+
+    resolve({
+      status: 'rejected',
+      reason: 'Aborted',
+    });
+
+    await jest.advanceTimersToNextTimerAsync();
+    expect(broker.cleanUpProvingJobState).toHaveBeenCalled();
+    expect(catchSpy).toHaveBeenCalledWith(new Error('Aborted'));
+  });
 });

yarn-project/prover-client/src/proving_broker/caching_broker_facade.ts

@@ -59,6 +59,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     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,160 +149,174 @@ 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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }
 
   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,
     );
   }

yarn-project/prover-client/src/proving_broker/proving_agent.test.ts

@@ -239,7 +239,7 @@ describe('ProvingAgent', () => {
     const inputs: ProvingJobInputs = { type: ProvingRequestType.BASE_PARITY, inputs: makeBaseParityInputs() };
     const job: ProvingJob = {
       id: randomBytes(8).toString('hex') as ProvingJobId,
-      blockNumber: 1,
+      epochNumber: 1,
       type: ProvingRequestType.BASE_PARITY,
       inputsUri: randomBytes(8).toString('hex') as ProofUri,
     };