refactor: use common DATA_DIRECTORY env var

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

@@ -18,7 +18,8 @@ export const ProverBrokerConfig = z.object({
   /** If starting a prover broker locally, the interval the broker checks for timed out jobs */
   proverBrokerPollIntervalMs: z.number(),
   /** If starting a prover broker locally, the directory to store broker data */
-  proverBrokerDataDirectory: z.string().optional(),
+  dataDirectory: z.string().optional(),
+  dataStoreMapSizeKB: z.number(),
 });
 
 export type ProverBrokerConfig = z.infer<typeof ProverBrokerConfig>;
@@ -39,9 +40,14 @@ export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig>
     description: 'If starting a prover broker locally, the max number of retries per proving job',
     ...numberConfigHelper(3),
   },
-  proverBrokerDataDirectory: {
-    env: 'PROVER_BROKER_DATA_DIRECTORY',
-    description: 'If starting a prover broker locally, the directory to store broker data',
+  dataDirectory: {
+    env: 'DATA_DIRECTORY',
+    description: 'Optional dir to store data. If omitted will store in memory.',
+  },
+  dataStoreMapSizeKB: {
+    env: 'DATA_STORE_MAP_SIZE_KB',
+    description: 'DB mapping size to be applied to all key/value stores',
+    ...numberConfigHelper(128 * 1_024 * 1_024), // Defaulted to 128 GB
   },
 };
 

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

@@ -25,9 +25,8 @@ export type ProverConfig = ActualProverConfig & {
   nodeUrl?: string;
   /** Identifier of the prover */
   proverId: Fr;
-  /** Where to store temporary data */
-  cacheDir?: string;
-
+  dataDirectory?: string;
+  dataStoreMapSizeKB: number;
   proverAgentCount: number;
 };
 
@@ -36,7 +35,8 @@ export const ProverConfigSchema = z.object({
   realProofs: z.boolean(),
   proverId: schemas.Fr,
   proverTestDelayMs: z.number(),
-  cacheDir: z.string().optional(),
+  dataDirectory: z.string().optional(),
+  dataStoreMapSizeKB: z.number(),
   proverAgentCount: z.number(),
 }) satisfies ZodFor<ProverConfig>;
 
@@ -61,16 +61,20 @@ export const proverConfigMappings: ConfigMappingsType<ProverConfig> = {
     description: 'Artificial delay to introduce to all operations to the test prover.',
     ...numberConfigHelper(0),
   },
-  cacheDir: {
-    env: 'PROVER_CACHE_DIR',
-    description: 'Where to store cache data generated while proving',
-    defaultValue: '/tmp/aztec-prover',
-  },
   proverAgentCount: {
     env: 'PROVER_AGENT_COUNT',
     description: 'The number of prover agents to start',
     ...numberConfigHelper(1),
   },
+  dataDirectory: {
+    env: 'DATA_DIRECTORY',
+    description: 'Optional dir to store data. If omitted will store in memory.',
+  },
+  dataStoreMapSizeKB: {
+    env: 'DATA_STORE_MAP_SIZE_KB',
+    description: 'DB mapping size to be applied to all key/value stores',
+    ...numberConfigHelper(128 * 1_024 * 1_024), // Defaulted to 128 GB
+  },
 };
 
 function parseProverId(str: string) {

yarn-project/foundation/src/config/env_var.ts

@@ -107,7 +107,6 @@ export type EnvVar =
   | 'PROVER_BROKER_JOB_TIMEOUT_MS'
   | 'PROVER_BROKER_POLL_INTERVAL_MS'
   | 'PROVER_BROKER_JOB_MAX_RETRIES'
-  | 'PROVER_BROKER_DATA_DIRECTORY'
   | 'PROVER_COORDINATION_NODE_URL'
   | 'PROVER_DISABLED'
   | 'PROVER_ID'
@@ -120,7 +119,6 @@ export type EnvVar =
   | 'PROVER_REAL_PROOFS'
   | 'PROVER_REQUIRED_CONFIRMATIONS'
   | 'PROVER_TEST_DELAY_MS'
-  | 'PROVER_CACHE_DIR'
   | 'PXE_BLOCK_POLLING_INTERVAL_MS'
   | 'PXE_L2_STARTING_BLOCK'
   | 'PXE_PROVER_ENABLED'

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

@@ -1,14 +1,22 @@
 import { type ProverBrokerConfig } from '@aztec/circuit-types';
 import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
 
+import { join } from 'path';
+
 import { ProvingBroker } from './proving_broker.js';
+import { ProvingBrokerDatabase } from './proving_broker_database.js';
 import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js';
 import { KVBrokerDatabase } from './proving_broker_database/persisted.js';
 
 export async function createAndStartProvingBroker(config: ProverBrokerConfig): Promise<ProvingBroker> {
-  const database = config.proverBrokerDataDirectory
-    ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory))
-    : new InMemoryBrokerDatabase();
+  let database: ProvingBrokerDatabase;
+  if (config.dataDirectory) {
+    const dataDir = join(config.dataDirectory, 'prover_broker');
+    const store = AztecLmdbStore.open(dataDir, config.dataStoreMapSizeKB);
+    database = new KVBrokerDatabase(store);
+  } else {
+    database = new InMemoryBrokerDatabase();
+  }
 
   const broker = new ProvingBroker(database, {
     jobTimeoutMs: config.proverBrokerJobTimeoutMs,

yarn-project/prover-client/src/tx-prover/tx-prover.ts

@@ -16,8 +16,6 @@ import { createDebugLogger } from '@aztec/foundation/log';
 import { NativeACVMSimulator } from '@aztec/simulator';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 
-import { join } from 'path';
-
 import { type ProverClientConfig } from '../config.js';
 import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
 import { CachingBrokerFacade } from '../proving_broker/caching_broker_facade.js';
@@ -33,8 +31,6 @@ export class TxProver implements EpochProverManager {
   private running = false;
   private agents: ProvingAgent[] = [];
 
-  private cacheDir?: string;
-
   private constructor(
     private config: ProverClientConfig,
     private telemetry: TelemetryClient,
@@ -44,7 +40,6 @@ export class TxProver implements EpochProverManager {
   ) {
     // TODO(palla/prover-node): Cache the paddingTx here, and not in each proving orchestrator,
     // so it can be reused across multiple ones and not recomputed every time.
-    this.cacheDir = this.config.cacheDir ? join(this.config.cacheDir, `tx_prover_${this.config.proverId}`) : undefined;
   }
 
   public createEpochProver(db: MerkleTreeWriteOperations, cache: ProverCache = new InMemoryProverCache()): EpochProver {

yarn-project/prover-node/src/factory.ts

@@ -77,8 +77,8 @@ export async function createProverNode(
   const walletClient = publisher.getClient();
   const bondManager = await createBondManager(rollupContract, walletClient, config);
 
-  const cacheDir = config.cacheDir ? join(config.cacheDir, `prover_${config.proverId}`) : undefined;
-  const cacheManager = new ProverCacheManager(cacheDir);
+  const cacheDir = config.dataDirectory ? join(config.dataDirectory, `prover_${config.proverId}`) : undefined;
+  const cacheManager = new ProverCacheManager(cacheDir, config.dataStoreMapSizeKB);
 
   return new ProverNode(
     prover,

yarn-project/prover-node/src/prover-cache/cache_manager.ts

@@ -14,26 +14,29 @@ const EPOCH_DIR_SEPARATOR = '_';
 const EPOCH_HASH_FILENAME = 'epoch_hash.txt';
 
 export class ProverCacheManager {
-  constructor(private cacheDir?: string, private log = createDebugLogger('aztec:prover-node:cache-manager')) {}
+  constructor(
+    private dataRootDir?: string,
+    private cacheMapSize?: number,
+    private log = createDebugLogger('aztec:prover-node:cache-manager'),
+  ) {}
 
   public async openCache(epochNumber: bigint, epochHash: Buffer): Promise<ProverCache> {
-    if (!this.cacheDir) {
+    if (!this.dataRootDir) {
       return new InMemoryProverCache();
     }
 
-    const epochDir = EPOCH_DIR_PREFIX + EPOCH_DIR_SEPARATOR + epochNumber;
-    const dataDir = join(this.cacheDir, epochDir);
+    const epochDataDir = join(this.dataRootDir, EPOCH_DIR_PREFIX + EPOCH_DIR_SEPARATOR + epochNumber);
 
-    const storedEpochHash = await readFile(join(dataDir, EPOCH_HASH_FILENAME), 'hex').catch(() => Buffer.alloc(0));
+    const storedEpochHash = await readFile(join(epochDataDir, EPOCH_HASH_FILENAME), 'hex').catch(() => Buffer.alloc(0));
     if (storedEpochHash.toString() !== epochHash.toString()) {
-      await rm(dataDir, { recursive: true, force: true });
+      await rm(epochDataDir, { recursive: true, force: true });
     }
 
-    await mkdir(dataDir, { recursive: true });
-    await writeFile(join(dataDir, EPOCH_HASH_FILENAME), epochHash.toString('hex'));
+    await mkdir(epochDataDir, { recursive: true });
+    await writeFile(join(epochDataDir, EPOCH_HASH_FILENAME), epochHash.toString('hex'));
 
-    const store = AztecLmdbStore.open(dataDir);
-    this.log.debug(`Created new database for epoch ${epochNumber} at ${dataDir}`);
+    const store = AztecLmdbStore.open(epochDataDir, this.cacheMapSize);
+    this.log.debug(`Created new database for epoch ${epochNumber} at ${epochDataDir}`);
     const cleanup = () => store.close();
     return new KVProverCache(store, cleanup);
   }
@@ -43,11 +46,11 @@ export class ProverCacheManager {
    * @param upToAndIncludingEpoch - The epoch number up to which to remove caches
    */
   public async removeStaleCaches(upToAndIncludingEpoch: bigint): Promise<void> {
-    if (!this.cacheDir) {
+    if (!this.dataRootDir) {
       return;
     }
 
-    const entries: Dirent[] = await readdir(this.cacheDir, { withFileTypes: true }).catch(() => []);
+    const entries: Dirent[] = await readdir(this.dataRootDir, { withFileTypes: true }).catch(() => []);
 
     for (const item of entries) {
       if (!item.isDirectory()) {
@@ -61,8 +64,10 @@ export class ProverCacheManager {
 
       const epochNumberInt = BigInt(epochNumber);
       if (epochNumberInt <= upToAndIncludingEpoch) {
-        this.log.info(`Removing old epoch database for epoch ${epochNumberInt} at ${join(this.cacheDir, item.name)}`);
-        await rm(join(this.cacheDir, item.name), { recursive: true });
+        this.log.info(
+          `Removing old epoch database for epoch ${epochNumberInt} at ${join(this.dataRootDir, item.name)}`,
+        );
+        await rm(join(this.dataRootDir, item.name), { recursive: true });
       }
     }
   }