fix: Split stores per component and split merkle tree operations (#8299)

This PR packs two changesets: 1. Spins off a `MerkleTreeAdminOperations` from `MerkleTreeOperations`, to have an interface that does not expose methods that commit changes to the underlying store. World state now exposes a method to get an `ephemeralFork` without these operations, that can be used for answering public simulation requests. Note that we do not yet enforce that no changes go into the store, that requires more changes to have "readonly" versions of the trees and store all the way down. 2. Moves creation of the underlying data store to the factory of each component, so each creates a new db as needed, instead of sharing a single one. This allows to keep separate db files for p2p, archive, and world state. As a bonus, it makes forking world state cheaper.
AztecProtocol · Sep 4, 2024 · 4ee69ac · 4ee69ac
1 parent 157dd11
commit 4ee69ac
Show file tree

Hide file tree

Showing 25 changed files with 187 additions and 180 deletions.
diff --git a/yarn-project/archiver/src/factory.ts b/yarn-project/archiver/src/factory.ts
@@ -1,4 +1,5 @@
-import { type AztecKVStore } from '@aztec/kv-store';
+import { createDebugLogger } from '@aztec/foundation/log';
+import { createStore } from '@aztec/kv-store/utils';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
@@ -7,14 +8,13 @@ import { type ArchiverConfig } from './archiver/config.js';
 import { KVArchiverDataStore } from './archiver/index.js';
 import { createArchiverClient } from './rpc/archiver_client.js';
 
-export function createArchiver(
+export async function createArchiver(
   config: ArchiverConfig,
-  store: AztecKVStore,
   telemetry: TelemetryClient = new NoopTelemetryClient(),
   opts: { blockUntilSync: boolean } = { blockUntilSync: true },
 ) {
   if (!config.archiverUrl) {
-    // first create and sync the archiver
+    const store = await createStore('archiver', config, createDebugLogger('aztec:archiver:lmdb'));
     const archiverStore = new KVArchiverDataStore(store, config.maxLogs);
     return Archiver.createAndSync(config, archiverStore, telemetry, opts.blockUntilSync);
   } else {

diff --git a/yarn-project/aztec-node/src/aztec-node/server.test.ts b/yarn-project/aztec-node/src/aztec-node/server.test.ts
@@ -4,12 +4,11 @@ import {
   type L1ToL2MessageSource,
   type L2BlockSource,
   type L2LogsSource,
+  type MerkleTreeAdminOperations,
   MerkleTreeId,
-  type MerkleTreeOperations,
   mockTxForRollup,
 } from '@aztec/circuit-types';
 import { AztecAddress, EthAddress, Fr, GasFees, GlobalVariables, MaxBlockNumber } from '@aztec/circuits.js';
-import { type AztecLmdbStore } from '@aztec/kv-store/lmdb';
 import { type P2P } from '@aztec/p2p';
 import { type GlobalVariableBuilder } from '@aztec/sequencer-client';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
@@ -24,7 +23,7 @@ import { AztecNodeService } from './server.js';
 describe('aztec node', () => {
   let p2p: MockProxy<P2P>;
   let globalVariablesBuilder: MockProxy<GlobalVariableBuilder>;
-  let merkleTreeOps: MockProxy<MerkleTreeOperations>;
+  let merkleTreeOps: MockProxy<MerkleTreeAdminOperations>;
 
   let lastBlockNumber: number;
 
@@ -42,7 +41,7 @@ describe('aztec node', () => {
     p2p = mock<P2P>();
 
     globalVariablesBuilder = mock<GlobalVariableBuilder>();
-    merkleTreeOps = mock<MerkleTreeOperations>();
+    merkleTreeOps = mock<MerkleTreeAdminOperations>();
 
     const worldState = mock<WorldStateSynchronizer>({
       getLatest: () => merkleTreeOps,
@@ -59,8 +58,6 @@ describe('aztec node', () => {
     // all txs use the same allowed FPC class
     const contractSource = mock<ContractDataSource>();
 
-    const store = mock<AztecLmdbStore>();
-
     const aztecNodeConfig: AztecNodeConfig = getConfigEnvVars();
 
     node = new AztecNodeService(
@@ -86,7 +83,6 @@ describe('aztec node', () => {
       31337,
       1,
       globalVariablesBuilder,
-      store,
       new TestCircuitVerifier(),
       new NoopTelemetryClient(),
     );

diff --git a/yarn-project/aztec-node/src/aztec-node/server.ts b/yarn-project/aztec-node/src/aztec-node/server.ts
@@ -50,10 +50,9 @@ import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { padArrayEnd } from '@aztec/foundation/collection';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
-import { type AztecKVStore } from '@aztec/kv-store';
-import { createStore, openTmpStore } from '@aztec/kv-store/utils';
+import { openTmpStore } from '@aztec/kv-store/utils';
 import { SHA256Trunc, StandardTree, UnbalancedTree } from '@aztec/merkle-tree';
-import { AztecKVTxPool, InMemoryAttestationPool, type P2P, createP2PClient } from '@aztec/p2p';
+import { InMemoryAttestationPool, type P2P, createP2PClient } from '@aztec/p2p';
 import { getCanonicalClassRegisterer } from '@aztec/protocol-contracts/class-registerer';
 import { getCanonicalFeeJuice } from '@aztec/protocol-contracts/fee-juice';
 import { getCanonicalInstanceDeployer } from '@aztec/protocol-contracts/instance-deployer';
@@ -77,7 +76,7 @@ import {
   type ProtocolContractAddresses,
 } from '@aztec/types/contracts';
 import { createValidatorClient } from '@aztec/validator-client';
-import { MerkleTrees, type WorldStateSynchronizer, createWorldStateSynchronizer } from '@aztec/world-state';
+import { type WorldStateSynchronizer, createWorldStateSynchronizer } from '@aztec/world-state';
 
 import { type AztecNodeConfig, getPackageInfo } from './config.js';
 import { NodeMetrics } from './node_metrics.js';
@@ -104,7 +103,6 @@ export class AztecNodeService implements AztecNode {
     protected readonly l1ChainId: number,
     protected readonly version: number,
     protected readonly globalVariableBuilder: GlobalVariableBuilder,
-    protected readonly merkleTreesDb: AztecKVStore,
     private proofVerifier: ClientProtocolCircuitVerifier,
     private telemetry: TelemetryClient,
     private log = createDebugLogger('aztec:node'),
@@ -131,7 +129,6 @@ export class AztecNodeService implements AztecNode {
     config: AztecNodeConfig,
     telemetry?: TelemetryClient,
     log = createDebugLogger('aztec:node'),
-    storeLog = createDebugLogger('aztec:node:lmdb'),
   ): Promise<AztecNodeService> {
     telemetry ??= new NoopTelemetryClient();
     const ethereumChain = createEthereumChain(config.l1RpcUrl, config.l1ChainId);
@@ -142,25 +139,17 @@ export class AztecNodeService implements AztecNode {
       );
     }
 
-    const store = await createStore(config, config.l1Contracts.rollupAddress, storeLog);
-
-    const archiver = await createArchiver(config, store, telemetry, { blockUntilSync: true });
+    const archiver = await createArchiver(config, telemetry, { blockUntilSync: true });
 
     // we identify the P2P transaction protocol by using the rollup contract address.
     // this may well change in future
     config.transactionProtocol = `/aztec/tx/${config.l1Contracts.rollupAddress.toString()}`;
 
     // create the tx pool and the p2p client, which will need the l2 block source
-    const p2pClient = await createP2PClient(
-      config,
-      store,
-      new AztecKVTxPool(store, telemetry),
-      new InMemoryAttestationPool(),
-      archiver,
-    );
+    const p2pClient = await createP2PClient(config, new InMemoryAttestationPool(), archiver, telemetry);
 
     // now create the merkle trees and the world state synchronizer
-    const worldStateSynchronizer = await createWorldStateSynchronizer(config, store, archiver, telemetry);
+    const worldStateSynchronizer = await createWorldStateSynchronizer(config, archiver, telemetry);
 
     // start both and wait for them to sync from the block source
     await Promise.all([p2pClient.start(), worldStateSynchronizer.start()]);
@@ -199,7 +188,6 @@ export class AztecNodeService implements AztecNode {
       ethereumChain.chainInfo.id,
       config.version,
       new GlobalVariableBuilder(config),
-      store,
       proofVerifier,
       telemetry,
       log,
@@ -726,13 +714,8 @@ export class AztecNodeService implements AztecNode {
     );
     const prevHeader = (await this.blockSource.getBlock(-1))?.header;
 
-    // Instantiate merkle trees so uncommitted updates by this simulation are local to it.
-    // TODO we should be able to remove this after https://github.com/AztecProtocol/aztec-packages/issues/1869
-    // So simulation of public functions doesn't affect the merkle trees.
-    const merkleTrees = await MerkleTrees.new(this.merkleTreesDb, new NoopTelemetryClient(), this.log);
-
     const publicProcessorFactory = new PublicProcessorFactory(
-      merkleTrees.asLatest(),
+      await this.worldStateSynchronizer.ephemeralFork(),
       this.contractDataSource,
       new WASMSimulator(),
       this.telemetry,

diff --git a/yarn-project/aztec/src/cli/cmds/start_archiver.ts b/yarn-project/aztec/src/cli/cmds/start_archiver.ts
@@ -21,8 +21,7 @@ export const startArchiver = async (options: any, signalHandlers: (() => Promise
   const archiverConfig = extractRelevantOptions<ArchiverConfig>(options, archiverConfigMappings, 'archiver');
 
   const storeLog = createDebugLogger('aztec:archiver:lmdb');
-  const rollupAddress = archiverConfig.l1Contracts.rollupAddress;
-  const store = await createStore(archiverConfig, rollupAddress, storeLog);
+  const store = await createStore('archiver', archiverConfig, storeLog);
   const archiverStore = new KVArchiverDataStore(store, archiverConfig.maxLogs);
 
   const telemetry = await createAndStartTelemetryClient(getTelemetryClientConfig());

diff --git a/yarn-project/circuit-types/src/interfaces/merkle_tree_operations.ts b/yarn-project/circuit-types/src/interfaces/merkle_tree_operations.ts
@@ -84,9 +84,7 @@ type LeafTypes = {
 
 export type MerkleTreeLeafType<ID extends MerkleTreeId> = LeafTypes[ID];
 
-/**
- * Defines the interface for operations on a set of Merkle Trees.
- */
+/** Defines the interface for operations on a set of Merkle Trees. */
 export interface MerkleTreeOperations {
   /**
    * Appends leaves to a given tree.
@@ -203,7 +201,10 @@ export interface MerkleTreeOperations {
     leaves: Buffer[],
     subtreeHeight: number,
   ): Promise<BatchInsertionResult<TreeHeight, SubtreeSiblingPathHeight>>;
+}
 
+/** Operations on merkle trees world state that can modify the underlying store. */
+export interface MerkleTreeAdminOperations extends MerkleTreeOperations {
   /**
    * Handles a single L2 block (i.e. Inserts the new note hashes into the merkle tree).
    * @param block - The L2 block to handle.

diff --git a/yarn-project/end-to-end/src/e2e_block_building.test.ts b/yarn-project/end-to-end/src/e2e_block_building.test.ts
@@ -339,7 +339,7 @@ describe('e2e_block_building', () => {
     });
 
     // Regression for https://github.com/AztecProtocol/aztec-packages/issues/8306
-    it.skip('can simulate public txs while building a block', async () => {
+    it('can simulate public txs while building a block', async () => {
       ({
         teardown,
         pxe,
@@ -368,7 +368,7 @@ describe('e2e_block_building', () => {
       }
 
       logger.info('Waiting for txs to be mined');
-      await Promise.all(txs.map(tx => tx.wait()));
+      await Promise.all(txs.map(tx => tx.wait({ proven: false, timeout: 600 })));
     });
   });
 });

diff --git a/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts b/yarn-project/end-to-end/src/e2e_prover/e2e_prover_test.ts
@@ -17,7 +17,6 @@ import {
   deployL1Contract,
 } from '@aztec/aztec.js';
 import { BBCircuitVerifier } from '@aztec/bb-prover';
-import { createStore } from '@aztec/kv-store/utils';
 import { RollupAbi } from '@aztec/l1-artifacts';
 import { TokenContract } from '@aztec/noir-contracts.js';
 import { type ProverNode, type ProverNodeConfig, createProverNode } from '@aztec/prover-node';
@@ -225,11 +224,8 @@ export class FullProverTest {
     // Creating temp store and archiver for fully proven prover node
 
     this.logger.verbose('Starting archiver for new prover node');
-    const store = await createStore({ dataDirectory: undefined }, this.l1Contracts.l1ContractAddresses.rollupAddress);
-
     const archiver = await createArchiver(
       { ...this.context.aztecNodeConfig, dataDirectory: undefined },
-      store,
       new NoopTelemetryClient(),
       { blockUntilSync: true },
     );

diff --git a/yarn-project/end-to-end/src/fixtures/snapshot_manager.ts b/yarn-project/end-to-end/src/fixtures/snapshot_manager.ts
@@ -8,7 +8,6 @@ import {
   type CompleteAddress,
   type DebugLogger,
   type DeployL1Contracts,
-  type EthAddress,
   EthCheatCodes,
   Fr,
   GrumpkinScalar,
@@ -23,7 +22,6 @@ import { asyncMap } from '@aztec/foundation/async-map';
 import { type Logger, createDebugLogger } from '@aztec/foundation/log';
 import { makeBackoff, retry } from '@aztec/foundation/retry';
 import { resolver, reviver } from '@aztec/foundation/serialize';
-import { createStore } from '@aztec/kv-store/utils';
 import { type ProverNode, type ProverNodeConfig, createProverNode } from '@aztec/prover-node';
 import { type PXEService, createPXEService, getPXEServiceConfig } from '@aztec/pxe';
 import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
@@ -235,21 +233,13 @@ async function teardown(context: SubsystemsContext | undefined) {
 }
 
 export async function createAndSyncProverNode(
-  rollupAddress: EthAddress,
   proverNodePrivateKey: `0x${string}`,
   aztecNodeConfig: AztecNodeConfig,
   aztecNode: AztecNode,
 ) {
   // Creating temp store and archiver for simulated prover node
-
-  const store = await createStore({ dataDirectory: undefined }, rollupAddress);
-
-  const archiver = await createArchiver(
-    { ...aztecNodeConfig, dataDirectory: undefined },
-    store,
-    new NoopTelemetryClient(),
-    { blockUntilSync: true },
-  );
+  const archiverConfig = { ...aztecNodeConfig, dataDirectory: undefined };
+  const archiver = await createArchiver(archiverConfig, new NoopTelemetryClient(), { blockUntilSync: true });
 
   // Prover node config is for simulated proofs
   const proverConfig: ProverNodeConfig = {
@@ -335,7 +325,6 @@ async function setupFromFresh(
 
   logger.verbose('Creating and syncing a simulated prover node...');
   const proverNode = await createAndSyncProverNode(
-    deployL1ContractsValues.l1ContractAddresses.rollupAddress,
     `0x${proverNodePrivateKey!.toString('hex')}`,
     aztecNodeConfig,
     aztecNode,
@@ -433,12 +422,7 @@ async function setupFromState(statePath: string, logger: Logger): Promise<Subsys
   const proverNodePrivateKey = getPrivateKeyFromIndex(2);
 
   logger.verbose('Creating and syncing a simulated prover node...');
-  const proverNode = await createAndSyncProverNode(
-    aztecNodeConfig.l1Contracts.rollupAddress,
-    `0x${proverNodePrivateKey!}`,
-    aztecNodeConfig,
-    aztecNode,
-  );
+  const proverNode = await createAndSyncProverNode(`0x${proverNodePrivateKey!}`, aztecNodeConfig, aztecNode);
 
   logger.verbose('Creating pxe...');
   const pxeConfig = getPXEServiceConfig();

diff --git a/yarn-project/end-to-end/src/public-testnet/e2e_public_testnet_transfer.test.ts b/yarn-project/end-to-end/src/public-testnet/e2e_public_testnet_transfer.test.ts
@@ -44,12 +44,7 @@ describe(`deploys and transfers a private only token`, () => {
         ? `0x${proverNodePrivateKey?.toString('hex')}`
         : proverConfig.publisherPrivateKey;
 
-    proverNode = await createAndSyncProverNode(
-      config.l1Contracts.rollupAddress,
-      proverConfig.publisherPrivateKey,
-      config,
-      aztecNode,
-    );
+    proverNode = await createAndSyncProverNode(proverConfig.publisherPrivateKey, config, aztecNode);
   }, 600_000);
 
   afterEach(async () => {

diff --git a/yarn-project/kv-store/src/utils.ts b/yarn-project/kv-store/src/utils.ts
@@ -1,17 +1,25 @@
 import { type EthAddress } from '@aztec/foundation/eth-address';
 import { type Logger, createDebugLogger } from '@aztec/foundation/log';
 
+import { join } from 'path';
+
 import { type AztecKVStore } from './interfaces/store.js';
 import { AztecLmdbStore } from './lmdb/store.js';
 
-export function createStore(
-  config: { dataDirectory: string | undefined } | (string | undefined),
-  rollupAddress: EthAddress,
-  log: Logger = createDebugLogger('aztec:kv-store'),
-) {
-  const dataDirectory = typeof config === 'string' ? config : config?.dataDirectory;
-  log.info(dataDirectory ? `Creating data store at directory ${dataDirectory}` : 'Creating ephemeral data store');
-  return initStoreForRollup(AztecLmdbStore.open(dataDirectory, false), rollupAddress, log);
+export type DataStoreConfig = { dataDirectory: string | undefined; l1Contracts: { rollupAddress: EthAddress } };
+
+export function createStore(name: string, config: DataStoreConfig, log: Logger = createDebugLogger('aztec:kv-store')) {
+  let { dataDirectory } = config;
+  if (typeof dataDirectory !== 'undefined') {
+    dataDirectory = join(dataDirectory, name);
+  }
+
+  log.info(
+    dataDirectory
+      ? `Creating ${name} data store at directory ${dataDirectory}`
+      : `Creating ${name} ephemeral data store`,
+  );
+  return initStoreForRollup(AztecLmdbStore.open(dataDirectory, false), config.l1Contracts.rollupAddress, log);
 }
 
 /**
@@ -21,7 +29,7 @@ export function createStore(
  * @param rollupAddress - The ETH address of the rollup contract
  * @returns A promise that resolves when the store is cleared, or rejects if the rollup address does not match
  */
-export async function initStoreForRollup<T extends AztecKVStore>(
+async function initStoreForRollup<T extends AztecKVStore>(
   store: T,
   rollupAddress: EthAddress,
   log?: Logger,

diff --git a/yarn-project/p2p/src/client/index.ts b/yarn-project/p2p/src/client/index.ts
@@ -1,24 +1,31 @@
 import { type L2BlockSource } from '@aztec/circuit-types';
+import { createDebugLogger } from '@aztec/foundation/log';
 import { type AztecKVStore } from '@aztec/kv-store';
+import { type DataStoreConfig, createStore } from '@aztec/kv-store/utils';
+import { type TelemetryClient } from '@aztec/telemetry-client';
+import { NoopTelemetryClient } from '@aztec/telemetry-client/noop';
 
 import { type AttestationPool } from '../attestation_pool/attestation_pool.js';
 import { P2PClient } from '../client/p2p_client.js';
 import { type P2PConfig } from '../config.js';
 import { DiscV5Service } from '../service/discV5_service.js';
 import { DummyP2PService } from '../service/dummy_service.js';
 import { LibP2PService, createLibP2PPeerId } from '../service/index.js';
-import { type TxPool } from '../tx_pool/index.js';
+import { AztecKVTxPool, type TxPool } from '../tx_pool/index.js';
 import { getPublicIp, resolveAddressIfNecessary, splitAddressPort } from '../util.js';
 
 export * from './p2p_client.js';
 
 export const createP2PClient = async (
-  config: P2PConfig,
-  store: AztecKVStore,
-  txPool: TxPool,
+  config: P2PConfig & DataStoreConfig,
   attestationsPool: AttestationPool,
   l2BlockSource: L2BlockSource,
+  telemetry: TelemetryClient = new NoopTelemetryClient(),
+  deps: { txPool?: TxPool; store?: AztecKVStore } = {},
 ) => {
+  const store = deps.store ?? (await createStore('p2p', config, createDebugLogger('aztec:p2p:lmdb')));
+  const txPool = deps.txPool ?? new AztecKVTxPool(store, telemetry);
+
   let p2pService;
 
   if (config.p2pEnabled) {