From 8e2fa60a8b08bd2235e976fd1223476608e11d86 Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Mon, 25 Nov 2024 17:22:45 +0000 Subject: [PATCH] wip --- yarn-project/prover-client/src/index.ts | 2 +- .../caching_broker_facade.test.ts | 111 ++++++++++++++++++ .../proving_broker/caching_broker_facade.ts | 2 +- .../prover_cache/memory.ts} | 0 .../prover-client/src/tx-prover/tx-prover.ts | 2 +- 5 files changed, 114 insertions(+), 3 deletions(-) create mode 100644 yarn-project/prover-client/src/proving_broker/caching_broker_facade.test.ts rename yarn-project/prover-client/src/{orchestrator/orchestrator_cache.ts => proving_broker/prover_cache/memory.ts} (100%) diff --git a/yarn-project/prover-client/src/index.ts b/yarn-project/prover-client/src/index.ts index d17ef228f211..56f3430e2c6f 100644 --- a/yarn-project/prover-client/src/index.ts +++ b/yarn-project/prover-client/src/index.ts @@ -3,4 +3,4 @@ export { EpochProverManager } from '@aztec/circuit-types'; export * from './tx-prover/tx-prover.js'; export * from './config.js'; export * from './tx-prover/factory.js'; -export * from './orchestrator/orchestrator_cache.js'; +export * from './proving_broker/prover_cache/memory.js'; diff --git a/yarn-project/prover-client/src/proving_broker/caching_broker_facade.test.ts b/yarn-project/prover-client/src/proving_broker/caching_broker_facade.test.ts new file mode 100644 index 000000000000..fa357379d902 --- /dev/null +++ b/yarn-project/prover-client/src/proving_broker/caching_broker_facade.test.ts @@ -0,0 +1,111 @@ +import { type ProvingJobProducer, ProvingRequestType, makePublicInputsAndRecursiveProof } from '@aztec/circuit-types'; +import { RECURSIVE_PROOF_LENGTH, VerificationKeyData, makeRecursiveProof } from '@aztec/circuits.js'; +import { makeBaseParityInputs, makeParityPublicInputs } from '@aztec/circuits.js/testing'; +import { AbortError } from '@aztec/foundation/error'; +import { promiseWithResolvers } from '@aztec/foundation/promise'; + +import { jest } from '@jest/globals'; +import { type MockProxy, mock } from 'jest-mock-extended'; + +import { CachingBrokerFacade } from './caching_broker_facade.js'; +import { InlineProofStore } from './proof_store.js'; +import { InMemoryProverCache } from './prover_cache/memory.js'; + +describe('CachingBrokerFacade', () => { + let facade: CachingBrokerFacade; + let cache: InMemoryProverCache; + let proofStore: InlineProofStore; + let broker: MockProxy; + + beforeAll(() => { + jest.useFakeTimers(); + }); + + beforeEach(() => { + broker = mock({ + enqueueProvingJob: jest.fn(), + getProvingJobStatus: jest.fn(), + removeAndCancelProvingJob: jest.fn(), + waitForJobToSettle: jest.fn(), + }); + cache = new InMemoryProverCache(); + proofStore = new InlineProofStore(); + facade = new CachingBrokerFacade(broker, cache, proofStore); + }); + + it('marks job as in progress', async () => { + const controller = new AbortController(); + void facade.getBaseParityProof(makeBaseParityInputs(), controller.signal); + + await jest.advanceTimersToNextTimerAsync(); + + expect(broker.enqueueProvingJob).toHaveBeenCalled(); + const job = broker.enqueueProvingJob.mock.calls[0][0]; + + await expect(cache.getProvingJobStatus(job.id)).resolves.toEqual({ status: 'in-queue' }); + controller.abort(); + }); + + it('removes the cached value if a job fails to enqueue', async () => { + const { promise, reject } = promiseWithResolvers(); + broker.enqueueProvingJob.mockResolvedValue(promise); + + void facade.getBaseParityProof(makeBaseParityInputs()).catch(() => {}); + await jest.advanceTimersToNextTimerAsync(); + + const job = broker.enqueueProvingJob.mock.calls[0][0]; + await expect(cache.getProvingJobStatus(job.id)).resolves.toEqual({ status: 'in-queue' }); + + reject(new Error('Failed to enqueue job')); + + await jest.advanceTimersToNextTimerAsync(); + await expect(cache.getProvingJobStatus(job.id)).resolves.toEqual({ status: 'not-found' }); + }); + + it('awaits existing job if in progress', async () => { + const { promise, reject } = promiseWithResolvers(); + broker.enqueueProvingJob.mockResolvedValue(promise); + + const inputs = makeBaseParityInputs(); + void facade.getBaseParityProof(inputs).catch(() => {}); + await jest.advanceTimersToNextTimerAsync(); + expect(broker.enqueueProvingJob).toHaveBeenCalledTimes(1); + + void facade.getBaseParityProof(inputs).catch(() => {}); + await jest.advanceTimersToNextTimerAsync(); + expect(broker.enqueueProvingJob).toHaveBeenCalledTimes(1); + + reject(new AbortError('Job was cancelled')); + }); + + it('reuses already cached results', async () => { + const { promise, resolve } = promiseWithResolvers(); + broker.enqueueProvingJob.mockResolvedValue(Promise.resolve()); + broker.waitForJobToSettle.mockResolvedValue(promise); + + const inputs = makeBaseParityInputs(); + void facade.getBaseParityProof(inputs); + await jest.advanceTimersToNextTimerAsync(); + + expect(broker.enqueueProvingJob).toHaveBeenCalledTimes(1); + 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(); + await expect(cache.getProvingJobStatus(job.id)).resolves.toEqual({ status: 'fulfilled', value: outputUri }); + + await expect(facade.getBaseParityProof(inputs)).resolves.toEqual(result); + expect(broker.enqueueProvingJob).toHaveBeenCalledTimes(1); // job was only ever enqueued once + }); +}); diff --git a/yarn-project/prover-client/src/proving_broker/caching_broker_facade.ts b/yarn-project/prover-client/src/proving_broker/caching_broker_facade.ts index b055e04f8dd9..c6c0356709c2 100644 --- a/yarn-project/prover-client/src/proving_broker/caching_broker_facade.ts +++ b/yarn-project/prover-client/src/proving_broker/caching_broker_facade.ts @@ -36,8 +36,8 @@ import { sha256 } from '@aztec/foundation/crypto'; import { createDebugLogger } from '@aztec/foundation/log'; import { retryUntil } from '@aztec/foundation/retry'; -import { InMemoryProverCache } from '../orchestrator/orchestrator_cache.js'; import { InlineProofStore, type ProofStore } from './proof_store.js'; +import { InMemoryProverCache } from './prover_cache/memory.js'; /** * A facade around a job broker that generates stable job ids and caches results diff --git a/yarn-project/prover-client/src/orchestrator/orchestrator_cache.ts b/yarn-project/prover-client/src/proving_broker/prover_cache/memory.ts similarity index 100% rename from yarn-project/prover-client/src/orchestrator/orchestrator_cache.ts rename to yarn-project/prover-client/src/proving_broker/prover_cache/memory.ts diff --git a/yarn-project/prover-client/src/tx-prover/tx-prover.ts b/yarn-project/prover-client/src/tx-prover/tx-prover.ts index a66381528c33..121af71c8014 100644 --- a/yarn-project/prover-client/src/tx-prover/tx-prover.ts +++ b/yarn-project/prover-client/src/tx-prover/tx-prover.ts @@ -20,9 +20,9 @@ import { join } from 'path'; import { type ProverClientConfig } from '../config.js'; import { ProvingOrchestrator } from '../orchestrator/orchestrator.js'; -import { InMemoryProverCache } from '../orchestrator/orchestrator_cache.js'; import { CachingBrokerFacade } from '../proving_broker/caching_broker_facade.js'; import { InlineProofStore } from '../proving_broker/proof_store.js'; +import { InMemoryProverCache } from '../proving_broker/prover_cache/memory.js'; import { ProvingAgent } from '../proving_broker/proving_agent.js'; /**