fix(builder): persist contractsDB across blocks within a checkpoint

yarn-project/end-to-end/src/e2e_epochs/epochs_mbps.parallel.test.ts

@@ -9,6 +9,7 @@ import { isL1ToL2MessageReady } from '@aztec/aztec.js/messaging';
 import { waitForTx } from '@aztec/aztec.js/node';
 import { RollupContract } from '@aztec/ethereum/contracts';
 import type { Operator } from '@aztec/ethereum/deploy-aztec-l1-contracts';
+import { waitUntilL1Timestamp } from '@aztec/ethereum/l1-tx-utils';
 import { asyncMap } from '@aztec/foundation/async-map';
 import { CheckpointNumber, SlotNumber } from '@aztec/foundation/branded-types';
 import { times, timesAsync } from '@aztec/foundation/collection';
@@ -17,6 +18,8 @@ import { retryUntil } from '@aztec/foundation/retry';
 import { bufferToHex } from '@aztec/foundation/string';
 import { executeTimeout } from '@aztec/foundation/timer';
 import { TestContract } from '@aztec/noir-test-contracts.js/Test';
+import { getSlotAtTimestamp, getTimestampForSlot } from '@aztec/stdlib/epoch-helpers';
+import { GasFees } from '@aztec/stdlib/gas';
 import { TxStatus } from '@aztec/stdlib/tx';
 
 import { jest } from '@jest/globals';
@@ -67,6 +70,7 @@ describe('e2e_epochs/epochs_mbps', () => {
     maxTxsPerBlock?: number;
     buildCheckpointIfEmpty?: boolean;
     deployCrossChainContract?: boolean;
+    skipPushProposedBlocksToArchiver?: boolean;
   }) {
     const { syncChainTip = 'checkpointed', deployCrossChainContract = false, ...setupOpts } = opts;
 
@@ -493,4 +497,83 @@ describe('e2e_epochs/epochs_mbps', () => {
     const multiBlockCheckpoint = await assertMultipleBlocksPerSlot(2, logger);
     await waitForProvenCheckpoint(multiBlockCheckpoint);
   });
+
+  it('deploys a contract and calls it in separate blocks within a slot', async () => {
+    await setupTest({
+      syncChainTip: 'checkpointed',
+      minTxsPerBlock: 1,
+      maxTxsPerBlock: 1,
+    });
+
+    // Prepare deploy tx for a new TestContract. Get the instance address so we can construct the call tx.
+    const highPriority = new GasFees(100, 100);
+    const lowPriority = new GasFees(1, 1);
+
+    const deployMethod = TestContract.deploy(wallet);
+    const deployInstance = await deployMethod.getInstance();
+    logger.warn(`Will deploy TestContract at ${deployInstance.address}`);
+
+    // Register the contract on the PXE so we can prove the call interaction against it.
+    await wallet.registerContract(deployInstance, TestContract.artifact);
+    const deployedContract = TestContract.at(deployInstance.address, wallet);
+
+    // Pre-prove both txs before starting sequencers. This ensures both arrive in the pool
+    // at the same time, so the sequencer can sort by priority fee for correct ordering.
+    logger.warn(`Pre-proving deploy tx (high priority) and call tx (low priority)`);
+    const deployTx = await proveInteraction(wallet, deployMethod, {
+      from,
+      fee: { gasSettings: { maxPriorityFeesPerGas: highPriority } },
+    });
+    const callTx = await proveInteraction(wallet, deployedContract.methods.emit_nullifier_public(new Fr(42)), {
+      from,
+      fee: { gasSettings: { maxPriorityFeesPerGas: lowPriority } },
+    });
+    logger.warn(`Pre-proved both txs`);
+
+    // Start the sequencers
+    await Promise.all(nodes.map(n => n.getSequencer()!.start()));
+    logger.warn(`Started all sequencers`);
+
+    // Wait until one L1 slot before the start of the next L2 slot.
+    // This ensures both txs land in the pending pool right before the proposer starts building.
+    // REFACTOR: This should go into a shared "waitUntilNextSlotStartsBuilding" utility
+    const currentL1Block = await test.l1Client.getBlock({ blockTag: 'latest' });
+    const currentTimestamp = currentL1Block.timestamp;
+    const currentSlot = getSlotAtTimestamp(currentTimestamp, test.constants);
+    const nextSlot = SlotNumber(currentSlot + 1);
+    const nextSlotTimestamp = getTimestampForSlot(nextSlot, test.constants);
+    const targetTimestamp = nextSlotTimestamp - BigInt(test.L1_BLOCK_TIME_IN_S);
+    logger.warn(`Waiting until L1 timestamp ${targetTimestamp} (one L1 slot before L2 slot ${nextSlot})`, {
+      currentTimestamp,
+      currentSlot,
+      nextSlot,
+      nextSlotTimestamp,
+      targetTimestamp,
+    });
+    await waitUntilL1Timestamp(test.l1Client, targetTimestamp, undefined, test.L2_SLOT_DURATION_IN_S * 3);
+
+    // Send both pre-proved txs simultaneously, waiting for them to be checkpointed.
+    const timeout = test.L2_SLOT_DURATION_IN_S * 5;
+    logger.warn(`Sending both txs and waiting for checkpointed receipts`);
+    const [deployReceipt, callReceipt] = await executeTimeout(
+      () => Promise.all([deployTx.send({ wait: { timeout } }), callTx.send({ wait: { timeout } })]),
+      timeout * 1000,
+    );
+    logger.warn(`Both txs checkpointed`, {
+      deployBlock: deployReceipt.blockNumber,
+      callBlock: callReceipt.blockNumber,
+    });
+
+    // Both txs should succeed (send throws on revert). Deploy should be in an earlier block.
+    expect(deployReceipt.blockNumber).toBeLessThan(callReceipt.blockNumber!);
+
+    // Verify both blocks belong to the same checkpoint.
+    const [deployCheckpointedBlock] = await context.aztecNode.getCheckpointedBlocks(deployReceipt.blockNumber!, 1);
+    const [callCheckpointedBlock] = await context.aztecNode.getCheckpointedBlocks(callReceipt.blockNumber!, 1);
+    expect(deployCheckpointedBlock.checkpointNumber).toBe(callCheckpointedBlock.checkpointNumber);
+    logger.warn(`Both blocks in checkpoint ${deployCheckpointedBlock.checkpointNumber}`);
+
+    // Wait for the checkpoint to be proven.
+    await waitForProvenCheckpoint(deployCheckpointedBlock.checkpointNumber);
+  });
 });

yarn-project/validator-client/src/checkpoint_builder.test.ts

@@ -10,7 +10,7 @@ import { Fr } from '@aztec/foundation/curves/bn254';
 import { EthAddress } from '@aztec/foundation/eth-address';
 import { TestDateProvider } from '@aztec/foundation/timer';
 import type { LightweightCheckpointBuilder } from '@aztec/prover-client/light';
-import type { PublicProcessor } from '@aztec/simulator/server';
+import type { PublicContractsDB, PublicProcessor } from '@aztec/simulator/server';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { L2Block } from '@aztec/stdlib/block';
 import type { ContractDataSource } from '@aztec/stdlib/contract';
@@ -22,16 +22,22 @@ import {
   type PublicProcessorLimits,
   type PublicProcessorValidator,
 } from '@aztec/stdlib/interfaces/server';
-import type { CheckpointGlobalVariables, GlobalVariables, ProcessedTx, Tx } from '@aztec/stdlib/tx';
+import {
+  type CheckpointGlobalVariables,
+  type GlobalVariables,
+  type ProcessedTx,
+  type Tx,
+  TxHash,
+} from '@aztec/stdlib/tx';
 import type { TelemetryClient } from '@aztec/telemetry-client';
 
-import { describe, expect, it } from '@jest/globals';
+import { describe, expect, it, jest } from '@jest/globals';
 import { type MockProxy, mock } from 'jest-mock-extended';
 
 import { CheckpointBuilder } from './checkpoint_builder.js';
 
 describe('CheckpointBuilder', () => {
-  let checkpointBuilder: CheckpointBuilder;
+  let checkpointBuilder: TestCheckpointBuilder;
   let lightweightCheckpointBuilder: MockProxy<LightweightCheckpointBuilder>;
   let fork: MockProxy<MerkleTreeWriteOperations>;
   let config: FullNodeBlockBuilderConfig;
@@ -56,6 +62,8 @@ describe('CheckpointBuilder', () => {
   };
 
   class TestCheckpointBuilder extends CheckpointBuilder {
+    declare public contractsDB: PublicContractsDB;
+
     public override makeBlockBuilderDeps(_globalVariables: GlobalVariables, _fork: MerkleTreeWriteOperations) {
       return Promise.resolve({ processor, validator });
     }
@@ -116,6 +124,50 @@ describe('CheckpointBuilder', () => {
     setupBuilder();
   });
 
+  describe('contractsDB checkpointing', () => {
+    let createCheckpointSpy: jest.SpiedFunction<() => void>;
+    let commitCheckpointSpy: jest.SpiedFunction<() => void>;
+    let revertCheckpointSpy: jest.SpiedFunction<() => void>;
+
+    beforeEach(() => {
+      const db = checkpointBuilder.contractsDB;
+      createCheckpointSpy = jest.spyOn(db, 'createCheckpoint');
+      commitCheckpointSpy = jest.spyOn(db, 'commitCheckpoint');
+      revertCheckpointSpy = jest.spyOn(db, 'revertCheckpoint');
+
+      lightweightCheckpointBuilder.getBlockCount.mockReturnValue(0);
+    });
+
+    async function mockSuccessfulBlock() {
+      const block = await L2Block.random(blockNumber);
+      lightweightCheckpointBuilder.addBlock.mockResolvedValue({ block, timings: {} });
+      processor.process.mockResolvedValue([[{ hash: TxHash.random() } as ProcessedTx], [], [], [], []]);
+      return block;
+    }
+
+    it('uses the same contractsDB across multiple block builds', async () => {
+      await mockSuccessfulBlock();
+      await checkpointBuilder.buildBlock([], blockNumber, 1000n);
+
+      await mockSuccessfulBlock();
+      await checkpointBuilder.buildBlock([], BlockNumber(blockNumber + 1), 1001n);
+
+      expect(createCheckpointSpy).toHaveBeenCalledTimes(2);
+      expect(commitCheckpointSpy).toHaveBeenCalledTimes(2);
+      expect(revertCheckpointSpy).not.toHaveBeenCalled();
+    });
+
+    it('calls revertCheckpoint when public processor fails', async () => {
+      processor.process.mockRejectedValue(new Error('processor failure'));
+
+      await expect(checkpointBuilder.buildBlock([], blockNumber, 1000n)).rejects.toThrow('processor failure');
+
+      expect(createCheckpointSpy).toHaveBeenCalledTimes(1);
+      expect(commitCheckpointSpy).not.toHaveBeenCalled();
+      expect(revertCheckpointSpy).toHaveBeenCalledTimes(1);
+    });
+  });
+
   describe('buildBlock', () => {
     it('builds a block successfully when transactions are processed', async () => {
       lightweightCheckpointBuilder.getBlockCount.mockReturnValue(0);

yarn-project/validator-client/src/checkpoint_builder.ts

@@ -45,6 +45,9 @@ export type { BuildBlockInCheckpointResult } from '@aztec/stdlib/interfaces/serv
 export class CheckpointBuilder implements ICheckpointBlockBuilder {
   private log: Logger;
 
+  /** Persistent contracts DB shared across all blocks in this checkpoint. */
+  protected contractsDB: PublicContractsDB;
+
   constructor(
     private checkpointBuilder: LightweightCheckpointBuilder,
     private fork: MerkleTreeWriteOperations,
@@ -59,6 +62,7 @@ export class CheckpointBuilder implements ICheckpointBlockBuilder {
       ...bindings,
       instanceId: `checkpoint-${checkpointBuilder.checkpointNumber}`,
     });
+    this.contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
   }
 
   getConstantData(): CheckpointGlobalVariables {
@@ -103,34 +107,44 @@ export class CheckpointBuilder implements ICheckpointBlockBuilder {
       ...this.capLimitsByCheckpointBudgets(opts),
     };
 
-    const [publicProcessorDuration, [processedTxs, failedTxs, usedTxs]] = await elapsed(() =>
-      processor.process(pendingTxs, cappedOpts, validator),
-    );
-
-    // Throw if we didn't collect a single valid tx and we're not allowed to build empty blocks
-    // (only the first block in a checkpoint can be empty)
-    if (processedTxs.length === 0 && this.checkpointBuilder.getBlockCount() > 0) {
-      throw new NoValidTxsError(failedTxs);
-    }
-
-    // Add block to checkpoint
-    const { block } = await this.checkpointBuilder.addBlock(globalVariables, processedTxs, {
-      expectedEndState: opts.expectedEndState,
-    });
+    // Create a block-level checkpoint on the contracts DB so we can roll back on failure
+    this.contractsDB.createCheckpoint();
 
-    this.log.debug('Built block within checkpoint', {
-      header: block.header.toInspect(),
-      processedTxs: processedTxs.map(tx => tx.hash.toString()),
-      failedTxs: failedTxs.map(tx => tx.tx.txHash.toString()),
-    });
+    try {
+      const [publicProcessorDuration, [processedTxs, failedTxs, usedTxs]] = await elapsed(() =>
+        processor.process(pendingTxs, cappedOpts, validator),
+      );
 
-    return {
-      block,
-      publicProcessorDuration,
-      numTxs: processedTxs.length,
-      failedTxs,
-      usedTxs,
-    };
+      // Throw if we didn't collect a single valid tx and we're not allowed to build empty blocks
+      // (only the first block in a checkpoint can be empty)
+      if (processedTxs.length === 0 && this.checkpointBuilder.getBlockCount() > 0) {
+        throw new NoValidTxsError(failedTxs);
+      }
+
+      // Add block to checkpoint
+      const { block } = await this.checkpointBuilder.addBlock(globalVariables, processedTxs, {
+        expectedEndState: opts.expectedEndState,
+      });
+
+      this.contractsDB.commitCheckpoint();
+
+      this.log.debug('Built block within checkpoint', {
+        header: block.header.toInspect(),
+        processedTxs: processedTxs.map(tx => tx.hash.toString()),
+        failedTxs: failedTxs.map(tx => tx.tx.txHash.toString()),
+      });
+
+      return {
+        block,
+        publicProcessorDuration,
+        numTxs: processedTxs.length,
+        failedTxs,
+        usedTxs,
+      };
+    } catch (err) {
+      this.contractsDB.revertCheckpoint();
+      throw err;
+    }
   }
 
   /** Completes the checkpoint and returns it. */
@@ -219,7 +233,7 @@ export class CheckpointBuilder implements ICheckpointBlockBuilder {
       ...(await getDefaultAllowedSetupFunctions()),
       ...(this.config.txPublicSetupAllowListExtend ?? []),
     ];
-    const contractsDB = new PublicContractsDB(this.contractDataSource, this.log.getBindings());
+    const contractsDB = this.contractsDB;
     const guardedFork = new GuardedMerkleTreeOperations(fork);
 
     const collectDebugLogs = this.debugLogStore.isEnabled;