feat: snapshotting for e2e p2p tests

yarn-project/end-to-end/Earthfile

@@ -2,7 +2,7 @@ VERSION 0.8
 
 e2e-p2p:
   LOCALLY
-  RUN ./scripts/e2e_test.sh ./src/e2e_p2p_network.test.ts
+  RUN ./scripts/e2e_test.sh ./src/e2e_p2p/ --runInBand
 
 e2e-l1-with-wall-time:
   LOCALLY

yarn-project/end-to-end/scripts/e2e_test.sh

@@ -1,13 +1,15 @@
 #!/bin/bash
 
-# Usage: ./e2e_test.sh <test>
+# Usage: ./e2e_test.sh <test> <...extra_args>
 # Optional environment variables:
 #   HARDWARE_CONCURRENCY (default: "")
 
 set -eu
 
 # Main positional parameter
 TEST="$1"
+shift
+
 # Default values for environment variables
 HARDWARE_CONCURRENCY="${HARDWARE_CONCURRENCY:-}"
 FAKE_PROOFS="${FAKE_PROOFS:-}"
@@ -18,4 +20,4 @@ if ! docker image ls --format '{{.Repository}}:{{.Tag}}' | grep -q "aztecprotoco
   exit 1
 fi
 
-docker run -e HARDWARE_CONCURRENCY="$HARDWARE_CONCURRENCY" -e FAKE_PROOFS="$FAKE_PROOFS" --rm aztecprotocol/end-to-end:$AZTEC_DOCKER_TAG "$TEST"
+docker run -e HARDWARE_CONCURRENCY="$HARDWARE_CONCURRENCY" -e FAKE_PROOFS="$FAKE_PROOFS" --rm aztecprotocol/end-to-end:$AZTEC_DOCKER_TAG "$TEST" $@

yarn-project/end-to-end/src/e2e_p2p/gossip_network.test.ts

@@ -0,0 +1,76 @@
+import { type AztecNodeService } from '@aztec/aztec-node';
+import { sleep } from '@aztec/aztec.js';
+
+import fs from 'fs';
+
+import { type NodeContext, createNodes } from '../fixtures/setup_p2p_test.js';
+import { P2PNetworkTest, WAIT_FOR_TX_TIMEOUT } from './p2p_network.js';
+import { createPXEServiceAndSubmitTransactions } from './shared.js';
+
+// Don't set this to a higher value than 9 because each node will use a different L1 publisher account and anvil seeds
+const NUM_NODES = 4;
+const NUM_TXS_PER_NODE = 2;
+const BOOT_NODE_UDP_PORT = 40600;
+
+const DATA_DIR = './data/gossip';
+
+describe('e2e_p2p_network', () => {
+  let t: P2PNetworkTest;
+  let nodes: AztecNodeService[];
+
+  beforeEach(async () => {
+    t = await P2PNetworkTest.create('e2e_p2p_network', NUM_NODES, BOOT_NODE_UDP_PORT);
+    await t.applyBaseSnapshots();
+    await t.setup();
+  });
+
+  afterEach(async () => {
+    await t.stopNodes(nodes);
+    await t.teardown();
+    for (let i = 0; i < NUM_NODES; i++) {
+      fs.rmSync(`${DATA_DIR}-${i}`, { recursive: true, force: true });
+    }
+  });
+
+  it('should rollup txs from all peers', async () => {
+    // create the bootstrap node for the network
+    if (!t.bootstrapNodeEnr) {
+      throw new Error('Bootstrap node ENR is not available');
+    }
+    // create our network of nodes and submit txs into each of them
+    // the number of txs per node and the number of txs per rollup
+    // should be set so that the only way for rollups to be built
+    // is if the txs are successfully gossiped around the nodes.
+    const contexts: NodeContext[] = [];
+    t.logger.info('Creating nodes');
+    nodes = await createNodes(
+      t.ctx.aztecNodeConfig,
+      t.peerIdPrivateKeys,
+      t.bootstrapNodeEnr,
+      NUM_NODES,
+      BOOT_NODE_UDP_PORT,
+      DATA_DIR,
+    );
+
+    // wait a bit for peers to discover each other
+    await sleep(4000);
+
+    t.logger.info('Submitting transactions');
+    for (const node of nodes) {
+      const context = await createPXEServiceAndSubmitTransactions(t.logger, node, NUM_TXS_PER_NODE);
+      contexts.push(context);
+    }
+
+    t.logger.info('Waiting for transactions to be mined');
+    // now ensure that all txs were successfully mined
+    await Promise.all(
+      contexts.flatMap((context, i) =>
+        context.txs.map(async (tx, j) => {
+          t.logger.info(`Waiting for tx ${i}-${j}: ${await tx.getTxHash()} to be mined`);
+          return tx.wait({ timeout: WAIT_FOR_TX_TIMEOUT });
+        }),
+      ),
+    );
+    t.logger.info('All transactions mined');
+  });
+});

yarn-project/end-to-end/src/e2e_p2p/p2p_network.ts

@@ -0,0 +1,156 @@
+import { type AztecNodeConfig, type AztecNodeService } from '@aztec/aztec-node';
+import { EthCheatCodes } from '@aztec/aztec.js';
+import { AZTEC_SLOT_DURATION, ETHEREUM_SLOT_DURATION, EthAddress } from '@aztec/circuits.js';
+import { type DebugLogger, createDebugLogger } from '@aztec/foundation/log';
+import { RollupAbi } from '@aztec/l1-artifacts';
+import { type BootstrapNode } from '@aztec/p2p';
+
+import getPort from 'get-port';
+import { getContract } from 'viem';
+import { privateKeyToAccount } from 'viem/accounts';
+
+import {
+  createBootstrapNodeFromPrivateKey,
+  createValidatorConfig,
+  generateNodePrivateKeys,
+  generatePeerIdPrivateKeys,
+} from '../fixtures/setup_p2p_test.js';
+import { type ISnapshotManager, type SubsystemsContext, createSnapshotManager } from '../fixtures/snapshot_manager.js';
+import { getPrivateKeyFromIndex } from '../fixtures/utils.js';
+
+// Use a fixed bootstrap node private key so that we can re-use the same snapshot and the nodes can find each other
+const BOOTSTRAP_NODE_PRIVATE_KEY = '080212208f988fc0899e4a73a5aee4d271a5f20670603a756ad8d84f2c94263a6427c591';
+export const WAIT_FOR_TX_TIMEOUT = AZTEC_SLOT_DURATION * 3;
+
+export class P2PNetworkTest {
+  private snapshotManager: ISnapshotManager;
+  private baseAccount;
+
+  public logger: DebugLogger;
+
+  public ctx!: SubsystemsContext;
+  public nodePrivateKeys: `0x${string}`[] = [];
+  public peerIdPrivateKeys: string[] = [];
+
+  public bootstrapNodeEnr: string = '';
+
+  constructor(
+    testName: string,
+    public bootstrapNode: BootstrapNode,
+    public bootNodePort: number,
+    private numberOfNodes: number,
+    initialValidatorAddress: string,
+    initialValidatorConfig: AztecNodeConfig,
+  ) {
+    this.logger = createDebugLogger(`aztec:e2e_p2p:${testName}`);
+
+    // Set up the base account and node private keys for the initial network deployment
+    this.baseAccount = privateKeyToAccount(`0x${getPrivateKeyFromIndex(0)!.toString('hex')}`);
+    this.nodePrivateKeys = generateNodePrivateKeys(1, numberOfNodes);
+    this.peerIdPrivateKeys = generatePeerIdPrivateKeys(numberOfNodes);
+
+    this.bootstrapNodeEnr = bootstrapNode.getENR().encodeTxt();
+
+    const initialValidators = [EthAddress.fromString(initialValidatorAddress)];
+
+    this.snapshotManager = createSnapshotManager(`e2e_p2p_network/${testName}`, process.env.E2E_DATA_PATH, {
+      ...initialValidatorConfig,
+      l1BlockTime: ETHEREUM_SLOT_DURATION,
+      salt: 420,
+      initialValidators,
+    });
+  }
+
+  static async create(testName: string, numberOfNodes: number, basePort?: number) {
+    const port = basePort || (await getPort());
+
+    const bootstrapNode = await createBootstrapNodeFromPrivateKey(BOOTSTRAP_NODE_PRIVATE_KEY, port);
+    const bootstrapNodeEnr = bootstrapNode.getENR().encodeTxt();
+
+    const initialValidatorConfig = await createValidatorConfig({} as AztecNodeConfig, bootstrapNodeEnr);
+    const intiailValidatorAddress = privateKeyToAccount(initialValidatorConfig.publisherPrivateKey).address;
+
+    return new P2PNetworkTest(
+      testName,
+      bootstrapNode,
+      port,
+      numberOfNodes,
+      intiailValidatorAddress,
+      initialValidatorConfig,
+    );
+  }
+
+  async applyBaseSnapshots() {
+    await this.snapshotManager.snapshot('add-validators', async ({ deployL1ContractsValues, aztecNodeConfig }) => {
+      const rollup = getContract({
+        address: deployL1ContractsValues.l1ContractAddresses.rollupAddress.toString(),
+        abi: RollupAbi,
+        client: deployL1ContractsValues.walletClient,
+      });
+
+      const txHashes: `0x${string}`[] = [];
+      for (let i = 0; i < this.numberOfNodes; i++) {
+        const account = privateKeyToAccount(this.nodePrivateKeys[i]!);
+        const txHash = await rollup.write.addValidator([account.address]);
+        txHashes.push(txHash);
+        this.logger.debug(`Adding ${account.address} as validator`);
+      }
+
+      // Remove the setup validator
+      const initialValidatorAddress = privateKeyToAccount(`0x${getPrivateKeyFromIndex(0)!.toString('hex')}`).address;
+      const txHash = await rollup.write.removeValidator([initialValidatorAddress]);
+      txHashes.push(txHash);
+
+      // Wait for all the transactions adding validators to be mined
+      await Promise.all(
+        txHashes.map(txHash =>
+          deployL1ContractsValues.publicClient.waitForTransactionReceipt({
+            hash: txHash,
+          }),
+        ),
+      );
+
+      //@note   Now we jump ahead to the next epoch such that the validator committee is picked
+      //        INTERVAL MINING: If we are using anvil interval mining this will NOT progress the time!
+      //        Which means that the validator set will still be empty! So anyone can propose.
+      const slotsInEpoch = await rollup.read.EPOCH_DURATION();
+      const timestamp = await rollup.read.getTimestampForSlot([slotsInEpoch]);
+      const cheatCodes = new EthCheatCodes(aztecNodeConfig.l1RpcUrl);
+      try {
+        await cheatCodes.warp(Number(timestamp));
+      } catch (err) {
+        this.logger.debug('Warp failed, time already satisfied');
+      }
+
+      // Send and await a tx to make sure we mine a block for the warp to correctly progress.
+      await deployL1ContractsValues.publicClient.waitForTransactionReceipt({
+        hash: await deployL1ContractsValues.walletClient.sendTransaction({
+          to: this.baseAccount.address,
+          value: 1n,
+          account: this.baseAccount,
+        }),
+      });
+    });
+  }
+
+  async setup() {
+    this.ctx = await this.snapshotManager.setup();
+
+    // TODO(md): make it such that the test can set these up
+    this.ctx.aztecNodeConfig.minTxsPerBlock = 4;
+    this.ctx.aztecNodeConfig.maxTxsPerBlock = 4;
+  }
+
+  async stopNodes(nodes: AztecNodeService[]) {
+    this.logger.info('Stopping nodes');
+    for (const node of nodes) {
+      await node.stop();
+    }
+    await this.bootstrapNode.stop();
+    this.logger.info('Nodes stopped');
+  }
+
+  async teardown() {
+    await this.snapshotManager.teardown();
+  }
+}

yarn-project/end-to-end/src/e2e_p2p/rediscovery.test.ts

@@ -0,0 +1,94 @@
+import { type AztecNodeService } from '@aztec/aztec-node';
+import { sleep } from '@aztec/aztec.js';
+
+import fs from 'fs';
+
+import { type NodeContext, createNode, createNodes } from '../fixtures/setup_p2p_test.js';
+import { P2PNetworkTest, WAIT_FOR_TX_TIMEOUT } from './p2p_network.js';
+import { createPXEServiceAndSubmitTransactions } from './shared.js';
+
+// Don't set this to a higher value than 9 because each node will use a different L1 publisher account and anvil seeds
+const NUM_NODES = 4;
+const NUM_TXS_PER_NODE = 2;
+const BOOT_NODE_UDP_PORT = 40400;
+
+const DATA_DIR = './data/rediscovery';
+
+describe('e2e_p2p_rediscovery', () => {
+  let t: P2PNetworkTest;
+  let nodes: AztecNodeService[];
+
+  beforeEach(async () => {
+    t = await P2PNetworkTest.create('e2e_p2p_rediscovery', NUM_NODES, BOOT_NODE_UDP_PORT);
+    await t.applyBaseSnapshots();
+    await t.setup();
+  });
+
+  afterEach(async () => {
+    await t.stopNodes(nodes);
+    await t.teardown();
+    for (let i = 0; i < NUM_NODES; i++) {
+      fs.rmSync(`${DATA_DIR}-${i}`, { recursive: true, force: true });
+    }
+  });
+
+  it('should re-discover stored peers without bootstrap node', async () => {
+    const contexts: NodeContext[] = [];
+    nodes = await createNodes(
+      t.ctx.aztecNodeConfig,
+      t.peerIdPrivateKeys,
+      t.bootstrapNodeEnr,
+      NUM_NODES,
+      BOOT_NODE_UDP_PORT,
+      DATA_DIR,
+    );
+
+    // wait a bit for peers to discover each other
+    await sleep(3000);
+
+    // stop bootstrap node
+    await t.bootstrapNode.stop();
+
+    // create new nodes from datadir
+    const newNodes: AztecNodeService[] = [];
+
+    // stop all nodes
+    for (let i = 0; i < NUM_NODES; i++) {
+      const node = nodes[i];
+      await node.stop();
+      t.logger.info(`Node ${i} stopped`);
+      await sleep(1200);
+
+      const newNode = await createNode(
+        t.ctx.aztecNodeConfig,
+        t.peerIdPrivateKeys[i],
+        i + 1 + BOOT_NODE_UDP_PORT,
+        undefined,
+        i,
+        `${DATA_DIR}-${i}`,
+      );
+      t.logger.info(`Node ${i} restarted`);
+      newNodes.push(newNode);
+    }
+    nodes = newNodes;
+
+    // wait a bit for peers to discover each other
+    await sleep(2000);
+
+    for (const node of newNodes) {
+      const context = await createPXEServiceAndSubmitTransactions(t.logger, node, NUM_TXS_PER_NODE);
+      contexts.push(context);
+    }
+
+    // now ensure that all txs were successfully mined
+
+    await Promise.all(
+      contexts.flatMap((context, i) =>
+        context.txs.map(async (tx, j) => {
+          t.logger.info(`Waiting for tx ${i}-${j}: ${await tx.getTxHash()} to be mined`);
+          return tx.wait({ timeout: WAIT_FOR_TX_TIMEOUT });
+        }),
+      ),
+    );
+  });
+});