test: end to end test node & pxe persistence

.circleci/config.yml

@@ -758,6 +758,17 @@ jobs:
           name: "Test"
           command: cond_spot_run_compose end-to-end 4 ./scripts/docker-compose.yml TEST=e2e_cli.test.ts
 
+  e2e-persistence:
+    docker:
+      - image: aztecprotocol/alpine-build-image
+    resource_class: small
+    steps:
+      - *checkout
+      - *setup_env
+      - run:
+          name: "Test"
+          command: cond_spot_run_compose end-to-end 4 ./scripts/docker-compose-no-sandbox.yml TEST=e2e_persistence.test.ts
+
   e2e-p2p:
     docker:
       - image: aztecprotocol/alpine-build-image
@@ -1205,6 +1216,7 @@ workflows:
       - uniswap-trade-on-l1-from-l2: *e2e_test
       - integration-l1-publisher: *e2e_test
       - integration-archiver-l1-to-l2: *e2e_test
+      - e2e-persistence: *e2e_test
       - e2e-p2p: *e2e_test
       - e2e-browser: *e2e_test
       - e2e-card-game: *e2e_test
@@ -1241,6 +1253,7 @@ workflows:
             - uniswap-trade-on-l1-from-l2
             - integration-l1-publisher
             - integration-archiver-l1-to-l2
+            - e2e-persistence
             - e2e-p2p
             - e2e-browser
             - e2e-card-game

yarn-project/aztec-node/src/aztec-node/server.ts

@@ -285,6 +285,7 @@ export class AztecNodeService implements AztecNode {
     await this.p2pClient.stop();
     await this.worldStateSynchronizer.stop();
     await this.blockSource.stop();
+    await this.merkleTreesDb.close();
     this.log.info(`Stopped`);
   }
 

yarn-project/aztec.js/src/account_manager/deploy_account_sent_tx.ts

@@ -3,7 +3,7 @@ import { TxHash, TxReceipt } from '@aztec/types';
 
 import { Wallet } from '../account/index.js';
 import { DefaultWaitOpts, SentTx, WaitOpts } from '../contract/index.js';
-import { waitForAccountSynch } from './util.js';
+import { waitForAccountSynch } from '../utils/account.js';
 
 /** Extends a transaction receipt with a wallet instance for the newly deployed contract. */
 export type DeployAccountTxReceipt = FieldsOf<TxReceipt> & {

yarn-project/aztec.js/src/account_manager/index.ts

@@ -7,10 +7,10 @@ import { Salt } from '../account/index.js';
 import { AccountInterface } from '../account/interface.js';
 import { DefaultWaitOpts, DeployMethod, WaitOpts } from '../contract/index.js';
 import { ContractDeployer } from '../contract_deployer/index.js';
+import { waitForAccountSynch } from '../utils/account.js';
 import { generatePublicKey } from '../utils/index.js';
 import { AccountWalletWithPrivateKey } from '../wallet/index.js';
 import { DeployAccountSentTx } from './deploy_account_sent_tx.js';
-import { waitForAccountSynch } from './util.js';
 
 /**
  * Manages a user account. Provides methods for calculating the account's address, deploying the account contract,

yarn-project/aztec.js/src/index.ts

@@ -45,6 +45,7 @@ export {
   EthCheatCodes,
   computeAuthWitMessageHash,
   waitForPXE,
+  waitForAccountSynch,
 } from './utils/index.js';
 
 export { createPXEClient } from './pxe_client.js';

yarn-project/aztec.js/src/account_manager/util.ts → yarn-project/aztec.js/src/utils/account.ts

@@ -1,7 +1,7 @@
 import { retryUntil } from '@aztec/foundation/retry';
 import { CompleteAddress, PXE } from '@aztec/types';
 
-import { WaitOpts } from '../contract/index.js';
+import { DefaultWaitOpts, WaitOpts } from '../contract/index.js';
 
 /**
  * Waits for the account to finish synchronizing with the PXE Service.
@@ -12,7 +12,7 @@ import { WaitOpts } from '../contract/index.js';
 export async function waitForAccountSynch(
   pxe: PXE,
   address: CompleteAddress,
-  { interval, timeout }: WaitOpts,
+  { interval, timeout }: WaitOpts = DefaultWaitOpts,
 ): Promise<void> {
   const publicKey = address.publicKey.toString();
   await retryUntil(

yarn-project/aztec.js/src/utils/index.ts

@@ -6,3 +6,4 @@ export * from './abi_types.js';
 export * from './cheat_codes.js';
 export * from './authwit.js';
 export * from './pxe.js';
+export * from './account.js';

yarn-project/end-to-end/src/e2e_2_pxes.test.ts

@@ -52,7 +52,7 @@ describe('e2e_2_pxes', () => {
       pxe: pxeB,
       accounts: accounts,
       wallets: [walletB],
-    } = await setupPXEService(1, aztecNode!, undefined, true));
+    } = await setupPXEService(1, aztecNode!, {}, undefined, true));
     [userB] = accounts;
   }, 100_000);
 

yarn-project/end-to-end/src/e2e_persistence.test.ts

@@ -0,0 +1,184 @@
+import { getUnsafeSchnorrAccount, getUnsafeSchnorrWallet } from '@aztec/accounts/single_key';
+import { AccountWallet, waitForAccountSynch } from '@aztec/aztec.js';
+import { CompleteAddress, EthAddress, Fq, Fr } from '@aztec/circuits.js';
+import { DeployL1Contracts } from '@aztec/ethereum';
+import { EasyPrivateTokenContract } from '@aztec/noir-contracts/EasyPrivateToken';
+
+import { mkdtemp } from 'fs/promises';
+import { tmpdir } from 'os';
+import { join } from 'path';
+
+import { EndToEndContext, setup } from './fixtures/utils.js';
+
+describe('Aztec persistence', () => {
+  /**
+   * These tests check that the Aztec Node and PXE can be shutdown and restarted without losing data.
+   *
+   * There are four scenarios to check:
+   * 1. Node and PXE are started with an existing databases
+   * 2. PXE is started with an existing database and connects to a Node with an empty database
+   * 3. PXE is started with an empty database and connects to a Node with an existing database
+   * 4. PXE is started with an empty database and connects to a Node with an empty database
+   *
+   * All four scenarios use the same L1 state, which is deployed in the `beforeAll` hook.
+   */
+
+  // the test contract and account deploying it
+  let contractAddress: CompleteAddress;
+  let ownerPrivateKey: Fq;
+  let ownerAddress: CompleteAddress;
+
+  // a directory where data will be persisted by components
+  // passing this through to the Node or PXE will control whether they use persisted data or not
+  let dataDirectory: string;
+
+  // state that is persisted between tests
+  let deployL1ContractsValues: DeployL1Contracts;
+
+  let context: EndToEndContext;
+
+  // deploy L1 contracts, start initial node & PXE, deploy test contract & shutdown node and PXE
+  beforeAll(async () => {
+    dataDirectory = await mkdtemp(join(tmpdir(), 'aztec-node-'));
+
+    const initialContext = await setup(0, { dataDirectory }, { dataDirectory });
+    deployL1ContractsValues = initialContext.deployL1ContractsValues;
+
+    ownerPrivateKey = Fq.random();
+    const ownerWallet = await getUnsafeSchnorrAccount(initialContext.pxe, ownerPrivateKey, Fr.ZERO).waitDeploy();
+    ownerAddress = ownerWallet.getCompleteAddress();
+
+    const deployer = EasyPrivateTokenContract.deploy(ownerWallet, 1000n, ownerWallet.getAddress());
+    await deployer.simulate({});
+
+    const contract = await deployer.send().deployed();
+    contractAddress = contract.completeAddress;
+
+    await initialContext.teardown();
+  }, 100_000);
+
+  describe.each([
+    [
+      // ie we were shutdown and now starting back up. Initial sync should be ~instant
+      'when starting Node and PXE with existing databases',
+      () => setup(0, { dataDirectory, deployL1ContractsValues }, { dataDirectory }),
+      1000,
+    ],
+    [
+      // ie our PXE was restarted, data kept intact and now connects to a "new" Node. Initial synch will synch from scratch
+      'when starting a PXE with an existing database, connected to a Node with database synched from scratch',
+      () => setup(0, { deployL1ContractsValues }, { dataDirectory }),
+      10_000,
+    ],
+  ])('%s', (_, contextSetup, timeout) => {
+    let ownerWallet: AccountWallet;
+    let contract: EasyPrivateTokenContract;
+
+    beforeEach(async () => {
+      context = await contextSetup();
+      ownerWallet = await getUnsafeSchnorrWallet(context.pxe, ownerAddress.address, ownerPrivateKey);
+      contract = await EasyPrivateTokenContract.at(contractAddress.address, ownerWallet);
+    }, timeout);
+
+    afterEach(async () => {
+      await context.teardown();
+    });
+
+    it('correctly restores balances', async () => {
+      // test for >0 instead of exact value so test isn't dependent on run order
+      await expect(contract.methods.getBalance(ownerWallet.getAddress()).view()).resolves.toBeGreaterThan(0n);
+    });
+
+    it('tracks new notes for the owner', async () => {
+      const balance = await contract.methods.getBalance(ownerWallet.getAddress()).view();
+      await contract.methods.mint(1000n, ownerWallet.getAddress()).send().wait();
+      await expect(contract.methods.getBalance(ownerWallet.getAddress()).view()).resolves.toEqual(balance + 1000n);
+    });
+
+    it('allows transfers of tokens from owner', async () => {
+      const otherWallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
+
+      const initialOwnerBalance = await contract.methods.getBalance(ownerWallet.getAddress()).view();
+      await contract.methods.transfer(500n, ownerWallet.getAddress(), otherWallet.getAddress()).send().wait();
+      const [ownerBalance, targetBalance] = await Promise.all([
+        contract.methods.getBalance(ownerWallet.getAddress()).view(),
+        contract.methods.getBalance(otherWallet.getAddress()).view(),
+      ]);
+
+      expect(ownerBalance).toEqual(initialOwnerBalance - 500n);
+      expect(targetBalance).toEqual(500n);
+    });
+  });
+
+  describe.each([
+    [
+      // ie. I'm setting up a new full node, sync from scratch and restore wallets/notes
+      'when starting the Node and PXE with empty databases',
+      () => setup(0, { deployL1ContractsValues }, {}),
+      10_000,
+    ],
+    [
+      // ie. I'm setting up a new PXE, restore wallets/notes from a Node
+      'when starting a PXE with an empty database connected to a Node with an existing database',
+      () => setup(0, { dataDirectory, deployL1ContractsValues }, {}),
+      10_000,
+    ],
+  ])('%s', (_, contextSetup, timeout) => {
+    beforeEach(async () => {
+      context = await contextSetup();
+    }, timeout);
+    afterEach(async () => {
+      await context.teardown();
+    });
+
+    it('pxe does not have the owner account', async () => {
+      await expect(context.pxe.getRecipient(ownerAddress.address)).resolves.toBeUndefined();
+    });
+
+    it('the node has the contract', async () => {
+      await expect(context.aztecNode.getContractData(contractAddress.address)).resolves.toBeDefined();
+    });
+
+    it('pxe does not know of the deployed contract', async () => {
+      await context.pxe.registerRecipient(ownerAddress);
+
+      const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
+      const contract = await EasyPrivateTokenContract.at(contractAddress.address, wallet);
+      await expect(contract.methods.getBalance(ownerAddress.address).view()).rejects.toThrowError(/Unknown contract/);
+    });
+
+    it("pxe does not have owner's notes", async () => {
+      await context.pxe.addContracts([
+        {
+          artifact: EasyPrivateTokenContract.artifact,
+          completeAddress: contractAddress,
+          portalContract: EthAddress.ZERO,
+        },
+      ]);
+      await context.pxe.registerRecipient(ownerAddress);
+
+      const wallet = await getUnsafeSchnorrAccount(context.pxe, Fq.random(), Fr.ZERO).waitDeploy();
+      const contract = await EasyPrivateTokenContract.at(contractAddress.address, wallet);
+      await expect(contract.methods.getBalance(ownerAddress.address).view()).resolves.toEqual(0n);
+    });
+
+    it('pxe restores notes after registering the owner', async () => {
+      await context.pxe.addContracts([
+        {
+          artifact: EasyPrivateTokenContract.artifact,
+          completeAddress: contractAddress,
+          portalContract: EthAddress.ZERO,
+        },
+      ]);
+
+      await context.pxe.registerAccount(ownerPrivateKey, ownerAddress.partialAddress);
+      const ownerWallet = await getUnsafeSchnorrAccount(context.pxe, ownerPrivateKey, ownerAddress).getWallet();
+      const contract = await EasyPrivateTokenContract.at(contractAddress.address, ownerWallet);
+
+      await waitForAccountSynch(context.pxe, ownerAddress, { interval: 1, timeout: 10 });
+
+      // check that notes total more than 0 so that this test isn't dependent on run order
+      await expect(contract.methods.getBalance(ownerAddress.address).view()).resolves.toBeGreaterThan(0n);
+    });
+  });
+});

yarn-project/end-to-end/src/fixtures/utils.ts

@@ -33,7 +33,7 @@ import {
   RollupAbi,
   RollupBytecode,
 } from '@aztec/l1-artifacts';
-import { PXEService, createPXEService, getPXEServiceConfig } from '@aztec/pxe';
+import { PXEService, PXEServiceConfig, createPXEService, getPXEServiceConfig } from '@aztec/pxe';
 import { SequencerClient } from '@aztec/sequencer-client';
 
 import * as path from 'path';
@@ -108,6 +108,7 @@ export const setupL1Contracts = async (
  * Sets up Private eXecution Environment (PXE).
  * @param numberOfAccounts - The number of new accounts to be created once the PXE is initiated.
  * @param aztecNode - An instance of Aztec Node.
+ * @param opts - Partial configuration for the PXE service.
  * @param firstPrivKey - The private key of the first account to be created.
  * @param logger - The logger to be used.
  * @param useLogSuffix - Whether to add a randomly generated suffix to the PXE debug logs.
@@ -116,6 +117,7 @@ export const setupL1Contracts = async (
 export async function setupPXEService(
   numberOfAccounts: number,
   aztecNode: AztecNode,
+  opts: Partial<PXEServiceConfig> = {},
   logger = getLogger(),
   useLogSuffix = false,
 ): Promise<{
@@ -136,7 +138,7 @@ export async function setupPXEService(
    */
   logger: DebugLogger;
 }> {
-  const pxeServiceConfig = getPXEServiceConfig();
+  const pxeServiceConfig = { ...getPXEServiceConfig(), ...opts };
   const pxe = await createPXEService(aztecNode, pxeServiceConfig, useLogSuffix);
 
   const wallets = await createAccounts(pxe, numberOfAccounts);
@@ -215,7 +217,12 @@ async function setupWithRemoteEnvironment(
 }
 
 /** Options for the e2e tests setup */
-type SetupOptions = { /** State load */ stateLoad?: string } & Partial<AztecNodeConfig>;
+type SetupOptions = {
+  /** State load */
+  stateLoad?: string;
+  /** Previously deployed contracts on L1 */
+  deployL1ContractsValues?: DeployL1Contracts;
+} & Partial<AztecNodeConfig>;
 
 /** Context for an end-to-end test as returned by the `setup` function */
 export type EndToEndContext = {
@@ -247,8 +254,13 @@ export type EndToEndContext = {
  * Sets up the environment for the end-to-end tests.
  * @param numberOfAccounts - The number of new accounts to be created once the PXE is initiated.
  * @param opts - Options to pass to the node initialization and to the setup script.
+ * @param pxeOpts - Options to pass to the PXE initialization.
  */
-export async function setup(numberOfAccounts = 1, opts: SetupOptions = {}): Promise<EndToEndContext> {
+export async function setup(
+  numberOfAccounts = 1,
+  opts: SetupOptions = {},
+  pxeOpts: Partial<PXEServiceConfig> = {},
+): Promise<EndToEndContext> {
   const config = { ...getConfigEnvVars(), ...opts };
 
   // Enable logging metrics to a local file named after the test suite
@@ -264,15 +276,15 @@ export async function setup(numberOfAccounts = 1, opts: SetupOptions = {}): Prom
 
   const logger = getLogger();
   const hdAccount = mnemonicToAccount(MNEMONIC);
+  const privKeyRaw = hdAccount.getHdKey().privateKey;
+  const publisherPrivKey = privKeyRaw === null ? null : Buffer.from(privKeyRaw);
 
   if (PXE_URL) {
     // we are setting up against a remote environment, l1 contracts are assumed to already be deployed
     return await setupWithRemoteEnvironment(hdAccount, config, logger, numberOfAccounts);
   }
 
-  const deployL1ContractsValues = await setupL1Contracts(config.rpcUrl, hdAccount, logger);
-  const privKeyRaw = hdAccount.getHdKey().privateKey;
-  const publisherPrivKey = privKeyRaw === null ? null : Buffer.from(privKeyRaw);
+  const deployL1ContractsValues = opts.deployL1ContractsValues ?? await setupL1Contracts(config.rpcUrl, hdAccount, logger);
 
   config.publisherPrivateKey = `0x${publisherPrivKey!.toString('hex')}`;
   config.l1Contracts.rollupAddress = deployL1ContractsValues.l1ContractAddresses.rollupAddress;
@@ -286,7 +298,7 @@ export async function setup(numberOfAccounts = 1, opts: SetupOptions = {}): Prom
   const aztecNode = await AztecNodeService.createAndSync(config);
   const sequencer = aztecNode.getSequencer();
 
-  const { pxe, accounts, wallets } = await setupPXEService(numberOfAccounts, aztecNode!, logger);
+  const { pxe, accounts, wallets } = await setupPXEService(numberOfAccounts, aztecNode!, pxeOpts, logger);
 
   const cheatCodes = CheatCodes.create(config.rpcUrl, pxe!);
 

yarn-project/foundation/src/fifo/memory_fifo.ts

@@ -60,6 +60,7 @@ export class MemoryFifo<T> {
    */
   public put(item: T) {
     if (this.flushing) {
+      this.log.warn('Discarding item because queue is flushing');
       return;
     } else if (this.waiting.length) {
       this.waiting.shift()!(item);