@aztec/prover-node 0.0.1-commit.b655e406 → 0.0.1-commit.b9865e97

package/src/job/epoch-proving-job.ts

@@ -1,23 +1,25 @@
-import { NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP } from '@aztec/constants';
 import { asyncPool } from '@aztec/foundation/async-pool';
-import { padArrayEnd } from '@aztec/foundation/collection';
-import { Fr } from '@aztec/foundation/fields';
-import { createLogger } from '@aztec/foundation/log';
+import { BlockNumber, EpochNumber } from '@aztec/foundation/branded-types';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { Timer } from '@aztec/foundation/timer';
+import { AVM_MAX_CONCURRENT_SIMULATIONS } from '@aztec/native';
 import { getVKTreeRoot } from '@aztec/noir-protocol-circuits-types/vk-tree';
 import { protocolContractsHash } from '@aztec/protocol-contracts';
 import { buildFinalBlobChallenges } from '@aztec/prover-client/helpers';
 import type { PublicProcessor, PublicProcessorFactory } from '@aztec/simulator/server';
+import { PublicSimulatorConfig } from '@aztec/stdlib/avm';
 import type { L2Block, L2BlockSource } from '@aztec/stdlib/block';
+import type { Checkpoint } from '@aztec/stdlib/checkpoint';
 import {
   type EpochProver,
   type EpochProvingJobState,
   EpochProvingJobTerminalState,
   type ForkMerkleTreeOperations,
 } from '@aztec/stdlib/interfaces/server';
+import { appendL1ToL2MessagesToTree } from '@aztec/stdlib/messaging';
 import { CheckpointConstantData } from '@aztec/stdlib/rollup';
-import { MerkleTreeId } from '@aztec/stdlib/trees';
 import type { ProcessedTx, Tx } from '@aztec/stdlib/tx';
 import { Attributes, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
 
@@ -40,10 +42,11 @@ export type EpochProvingJobOptions = {
  */
 export class EpochProvingJob implements Traceable {
   private state: EpochProvingJobState = 'initialized';
-  private log = createLogger('prover-node:epoch-proving-job');
+  private log: Logger;
   private uuid: string;
 
   private runPromise: Promise<void> | undefined;
+  private abortController = new AbortController();
   private epochCheckPromise: RunningPromise | undefined;
   private deadlineTimeoutHandler: NodeJS.Timeout | undefined;
 
@@ -54,14 +57,19 @@ export class EpochProvingJob implements Traceable {
     private dbProvider: Pick<ForkMerkleTreeOperations, 'fork'>,
     private prover: EpochProver,
     private publicProcessorFactory: PublicProcessorFactory,
-    private publisher: Pick<ProverNodePublisher, 'submitEpochProof'>,
+    private publisher: Pick<ProverNodePublisher, 'submitEpochProof' | 'analyzeEpochProofSubmission'>,
     private l2BlockSource: L2BlockSource | undefined,
     private metrics: ProverNodeJobMetrics,
     private deadline: Date | undefined,
     private config: EpochProvingJobOptions,
+    bindings?: LoggerBindings,
   ) {
     validateEpochProvingJobData(data);
     this.uuid = crypto.randomUUID();
+    this.log = createLogger('prover-node:epoch-proving-job', {
+      ...bindings,
+      instanceId: `epoch-${data.epochNumber}`,
+    });
     this.tracer = metrics.tracer;
   }
 
@@ -73,7 +81,7 @@ export class EpochProvingJob implements Traceable {
     return this.state;
   }
 
-  public getEpochNumber(): bigint {
+  public getEpochNumber(): EpochNumber {
     return this.data.epochNumber;
   }
 
@@ -89,8 +97,8 @@ export class EpochProvingJob implements Traceable {
     return this.data.epochNumber;
   }
 
-  private get blocks() {
-    return this.data.blocks;
+  private get checkpoints() {
+    return this.data.checkpoints;
   }
 
   private get txs() {
@@ -105,7 +113,7 @@ export class EpochProvingJob implements Traceable {
    * Proves the given epoch and submits the proof to L1.
    */
   @trackSpan('EpochProvingJob.run', function () {
-    return { [Attributes.EPOCH_NUMBER]: Number(this.data.epochNumber) };
+    return { [Attributes.EPOCH_NUMBER]: this.data.epochNumber };
   })
   public async run() {
     this.scheduleDeadlineStop();
@@ -114,14 +122,22 @@ export class EpochProvingJob implements Traceable {
     }
 
     const attestations = this.attestations.map(attestation => attestation.toViem());
-    const epochNumber = Number(this.epochNumber);
-    const epochSizeBlocks = this.blocks.length;
-    const epochSizeTxs = this.blocks.reduce((total, current) => total + current.body.txEffects.length, 0);
-    const [fromBlock, toBlock] = [this.blocks[0].number, this.blocks.at(-1)!.number];
-    this.log.info(`Starting epoch ${epochNumber} proving job with blocks ${fromBlock} to ${toBlock}`, {
+    const epochNumber = this.epochNumber;
+    const epochSizeCheckpoints = this.checkpoints.length;
+    const epochSizeBlocks = this.checkpoints.reduce((accum, checkpoint) => accum + checkpoint.blocks.length, 0);
+    const epochSizeTxs = this.checkpoints.reduce(
+      (accum, checkpoint) =>
+        accum + checkpoint.blocks.reduce((accumC, block) => accumC + block.body.txEffects.length, 0),
+      0,
+    );
+    const fromCheckpoint = this.checkpoints[0].number;
+    const toCheckpoint = this.checkpoints.at(-1)!.number;
+    const fromBlock = this.checkpoints[0].blocks[0].number;
+    const toBlock = this.checkpoints.at(-1)!.blocks.at(-1)!.number;
+    this.log.info(`Starting epoch ${epochNumber} proving job with checkpoints ${fromCheckpoint} to ${toCheckpoint}`, {
       fromBlock,
       toBlock,
-      epochSizeBlocks,
+      epochSizeTxs,
       epochNumber,
       uuid: this.uuid,
     });
@@ -132,84 +148,121 @@ export class EpochProvingJob implements Traceable {
     this.runPromise = promise;
 
     try {
-      const blobFieldsPerCheckpoint = this.blocks.map(block => block.getCheckpointBlobFields());
+      const blobTimer = new Timer();
+      const blobFieldsPerCheckpoint = this.checkpoints.map(checkpoint => checkpoint.toBlobFields());
       const finalBlobBatchingChallenges = await buildFinalBlobChallenges(blobFieldsPerCheckpoint);
+      this.metrics.recordBlobProcessing(blobTimer.ms());
 
-      // TODO(#17027): Enable multiple blocks per checkpoint.
-      // Total number of checkpoints equals number of blocks because we currently build a checkpoint with only one block.
-      const totalNumCheckpoints = epochSizeBlocks;
-
-      this.prover.startNewEpoch(epochNumber, totalNumCheckpoints, finalBlobBatchingChallenges);
+      this.prover.startNewEpoch(epochNumber, epochSizeCheckpoints, finalBlobBatchingChallenges);
+      const chonkTimer = new Timer();
       await this.prover.startChonkVerifierCircuits(Array.from(this.txs.values()));
+      this.metrics.recordChonkVerifier(chonkTimer.ms());
 
-      await asyncPool(this.config.parallelBlockLimit ?? 32, this.blocks, async block => {
-        this.checkState();
+      // Everything in the epoch should have the same chainId and version.
+      const { chainId, version } = this.checkpoints[0].blocks[0].header.globalVariables;
 
-        const globalVariables = block.header.globalVariables;
-        const txs = this.getTxs(block);
-        const l1ToL2Messages = this.getL1ToL2Messages(block);
-        const previousHeader = this.getBlockHeader(block.number - 1)!;
-
-        this.log.verbose(`Starting processing block ${block.number}`, {
-          number: block.number,
-          blockHash: (await block.hash()).toString(),
-          lastArchive: block.header.lastArchive.root,
-          noteHashTreeRoot: block.header.state.partial.noteHashTree.root,
-          nullifierTreeRoot: block.header.state.partial.nullifierTree.root,
-          publicDataTreeRoot: block.header.state.partial.publicDataTree.root,
-          previousHeader: previousHeader.hash(),
-          uuid: this.uuid,
-          ...globalVariables,
-        });
+      const previousBlockHeaders = this.gatherPreviousBlockHeaders();
+
+      const allCheckpointsTimer = new Timer();
 
+      const parallelism = this.config.parallelBlockLimit
+        ? this.config.parallelBlockLimit
+        : AVM_MAX_CONCURRENT_SIMULATIONS > 0
+          ? AVM_MAX_CONCURRENT_SIMULATIONS
+          : this.checkpoints.length;
+
+      await this.processCheckpoints(parallelism, async checkpoint => {
+        this.checkState();
+        const checkpointTimer = new Timer();
+
+        const checkpointIndex = checkpoint.number - fromCheckpoint;
         const checkpointConstants = CheckpointConstantData.from({
-          chainId: globalVariables.chainId,
-          version: globalVariables.version,
+          chainId,
+          version,
           vkTreeRoot: getVKTreeRoot(),
           protocolContractsHash: protocolContractsHash,
           proverId: this.prover.getProverId().toField(),
-          slotNumber: globalVariables.slotNumber,
-          coinbase: globalVariables.coinbase,
-          feeRecipient: globalVariables.feeRecipient,
-          gasFees: globalVariables.gasFees,
+          slotNumber: checkpoint.header.slotNumber,
+          coinbase: checkpoint.header.coinbase,
+          feeRecipient: checkpoint.header.feeRecipient,
+          gasFees: checkpoint.header.gasFees,
+        });
+        const previousHeader = previousBlockHeaders[checkpointIndex];
+        const l1ToL2Messages = this.getL1ToL2Messages(checkpoint);
+
+        this.log.debug(`Starting processing checkpoint ${checkpoint.number}`, {
+          number: checkpoint.number,
+          checkpointHash: checkpoint.hash().toString(),
+          headerHash: checkpoint.header.hash().toString(),
+          numL1ToL2Messages: l1ToL2Messages.length,
+          previousBlockNumber: previousHeader.globalVariables.blockNumber,
+          uuid: this.uuid,
         });
 
-        // TODO(#17027): Enable multiple blocks per checkpoint.
-        // Each checkpoint has only one block.
-        const totalNumBlocks = 1;
-        const checkpointIndex = block.number - fromBlock;
         await this.prover.startNewCheckpoint(
           checkpointIndex,
           checkpointConstants,
           l1ToL2Messages,
-          totalNumBlocks,
-          blobFieldsPerCheckpoint[checkpointIndex].length,
+          checkpoint.blocks.length,
           previousHeader,
         );
 
-        // Start block proving
-        await this.prover.startNewBlock(block.number, globalVariables.timestamp, txs.length);
+        for (let blockIndex = 0; blockIndex < checkpoint.blocks.length; blockIndex++) {
+          const blockTimer = new Timer();
+          const block = checkpoint.blocks[blockIndex];
+          const globalVariables = block.header.globalVariables;
+          const txs = this.getTxs(block);
+
+          this.log.verbose(`Starting processing block ${block.number}`, {
+            number: block.number,
+            blockHash: (await block.hash()).toString(),
+            lastArchive: block.header.lastArchive.root,
+            noteHashTreeRoot: block.header.state.partial.noteHashTree.root,
+            nullifierTreeRoot: block.header.state.partial.nullifierTree.root,
+            publicDataTreeRoot: block.header.state.partial.publicDataTree.root,
+            ...globalVariables,
+            numTxs: txs.length,
+          });
 
-        // Process public fns
-        const db = await this.createFork(block.number - 1, l1ToL2Messages);
-        const publicProcessor = this.publicProcessorFactory.create(db, globalVariables, {
-          skipFeeEnforcement: true,
-          clientInitiatedSimulation: false,
-          proverId: this.prover.getProverId().toField(),
-        });
-        const processed = await this.processTxs(publicProcessor, txs);
-        await this.prover.addTxs(processed);
-        await db.close();
-        this.log.verbose(`Processed all ${txs.length} txs for block ${block.number}`, {
-          blockNumber: block.number,
-          blockHash: (await block.hash()).toString(),
-          uuid: this.uuid,
-        });
+          // Start block proving
+          await this.prover.startNewBlock(block.number, globalVariables.timestamp, txs.length);
+
+          // Process public fns. L1 to L2 messages are only inserted for the first block of a checkpoint,
+          // as the fork for subsequent blocks already includes them from the previous block's synced state.
+          {
+            await using db = await this.createFork(
+              BlockNumber(block.number - 1),
+              blockIndex === 0 ? l1ToL2Messages : undefined,
+            );
+            this.checkState();
+            const config = PublicSimulatorConfig.from({
+              proverId: this.prover.getProverId().toField(),
+              skipFeeEnforcement: false,
+              collectDebugLogs: false,
+              collectHints: true,
+              collectPublicInputs: true,
+              collectStatistics: false,
+            });
+            const publicProcessor = this.publicProcessorFactory.create(db, globalVariables, config);
+            const processed = await this.processTxs(publicProcessor, txs);
+            this.checkState();
+            await this.prover.addTxs(processed);
+          }
+          this.checkState();
+          this.log.verbose(`Processed all ${txs.length} txs for block ${block.number}`, {
+            blockNumber: block.number,
+            blockHash: (await block.hash()).toString(),
+            uuid: this.uuid,
+          });
 
-        // Mark block as completed to pad it
-        const expectedBlockHeader = block.getBlockHeader();
-        await this.prover.setBlockCompleted(block.number, expectedBlockHeader);
+          // Mark block as completed to pad it
+          const expectedBlockHeader = block.header;
+          await this.prover.setBlockCompleted(block.number, expectedBlockHeader);
+          this.metrics.recordBlockProcessing(blockTimer.ms());
+        }
+        this.metrics.recordCheckpointProcessing(checkpointTimer.ms());
       });
+      this.metrics.recordAllCheckpointsProcessing(allCheckpointsTimer.ms());
 
       const executionTime = timer.ms();
 
@@ -221,16 +274,29 @@ export class EpochProvingJob implements Traceable {
 
       if (this.config.skipSubmitProof) {
         this.log.info(
-          `Proof publishing is disabled. Dropping valid proof for epoch ${epochNumber} (blocks ${fromBlock} to ${toBlock})`,
+          `Proof publishing is disabled. Analyzing estimated L1 fees for epoch ${epochNumber} (checkpoints ${fromCheckpoint} to ${toCheckpoint})`,
         );
+        try {
+          await this.publisher.analyzeEpochProofSubmission({
+            fromCheckpoint,
+            toCheckpoint,
+            epochNumber,
+            publicInputs,
+            proof,
+            batchedBlobInputs,
+            attestations,
+          });
+        } catch (err) {
+          this.log.warn(`Failed to analyze estimated L1 fees for epoch ${epochNumber}`, err);
+        }
         this.state = 'completed';
-        this.metrics.recordProvingJob(executionTime, timer.ms(), epochSizeBlocks, epochSizeTxs);
+        this.metrics.recordProvingJob(executionTime, timer.ms(), epochSizeCheckpoints, epochSizeBlocks, epochSizeTxs);
         return;
       }
 
       const success = await this.publisher.submitEpochProof({
-        fromBlock,
-        toBlock,
+        fromCheckpoint,
+        toCheckpoint,
         epochNumber,
         publicInputs,
         proof,
@@ -241,12 +307,12 @@ export class EpochProvingJob implements Traceable {
         throw new Error('Failed to submit epoch proof to L1');
       }
 
-      this.log.info(`Submitted proof for epoch ${epochNumber} (blocks ${fromBlock} to ${toBlock})`, {
+      this.log.info(`Submitted proof for epoch ${epochNumber} (checkpoints ${fromCheckpoint} to ${toCheckpoint})`, {
         epochNumber,
         uuid: this.uuid,
       });
       this.state = 'completed';
-      this.metrics.recordProvingJob(executionTime, timer.ms(), epochSizeBlocks, epochSizeTxs);
+      this.metrics.recordProvingJob(executionTime, timer.ms(), epochSizeCheckpoints, epochSizeBlocks, epochSizeTxs);
     } catch (err: any) {
       if (err && err.name === 'HaltExecutionError') {
         this.log.warn(`Halted execution of epoch ${epochNumber} prover job`, {
@@ -269,25 +335,63 @@ export class EpochProvingJob implements Traceable {
   }
 
   /**
-   * Create a new db fork for tx processing, inserting all L1 to L2.
+   * Create a new db fork for tx processing, optionally inserting L1 to L2 messages.
+   * L1 to L2 messages should only be inserted for the first block in a checkpoint,
+   * as subsequent blocks' synced state already includes them.
    * REFACTOR: The prover already spawns a db fork of its own for each block, so we may be able to do away with just one fork.
    */
-  private async createFork(blockNumber: number, l1ToL2Messages: Fr[]) {
-    const db = await this.dbProvider.fork(blockNumber);
-    const l1ToL2MessagesPadded = padArrayEnd<Fr, number>(
-      l1ToL2Messages,
-      Fr.ZERO,
-      NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
-      'Too many L1 to L2 messages',
-    );
-    this.log.verbose(`Creating fork at ${blockNumber} with ${l1ToL2Messages.length} L1 to L2 messages`, {
-      blockNumber,
-      l1ToL2Messages: l1ToL2MessagesPadded.map(m => m.toString()),
-    });
-    await db.appendLeaves(MerkleTreeId.L1_TO_L2_MESSAGE_TREE, l1ToL2MessagesPadded);
+  private async createFork(blockNumber: BlockNumber, l1ToL2Messages: Fr[] | undefined) {
+    this.log.verbose(`Creating fork at ${blockNumber}`, { blockNumber });
+    // temporary stack to control fork lifetime
+    await using cleanup = new AsyncDisposableStack();
+    const db = cleanup.use(await this.dbProvider.fork(blockNumber));
+
+    if (l1ToL2Messages !== undefined) {
+      this.log.verbose(`Inserting ${l1ToL2Messages.length} L1 to L2 messages in fork`, {
+        blockNumber,
+        l1ToL2Messages: l1ToL2Messages.map(m => m.toString()),
+      });
+      await appendL1ToL2MessagesToTree(db, l1ToL2Messages);
+    }
+
+    // everything run succesfully so we can release this stack and give control of the fork's lifetime to the caller
+    cleanup.move();
     return db;
   }
 
+  private async processCheckpoints(
+    parallelism: number,
+    processCheckpoint: (checkpoint: Checkpoint) => Promise<void>,
+  ): Promise<void> {
+    let hasError = false;
+    let firstError: unknown;
+
+    await asyncPool(Math.max(parallelism, 1), this.checkpoints, async checkpoint => {
+      if (hasError || this.abortController.signal.aborted) {
+        return;
+      }
+
+      try {
+        this.checkState();
+        await processCheckpoint(checkpoint);
+      } catch (err) {
+        if (!hasError) {
+          hasError = true;
+          firstError = err;
+          this.failProcessing();
+        }
+      }
+    });
+
+    if (hasError) {
+      throw firstError;
+    }
+
+    if (this.abortController.signal.aborted) {
+      this.checkState();
+    }
+  }
+
   private progressState(state: EpochProvingJobState) {
     this.checkState();
     this.state = state;
@@ -301,12 +405,24 @@ export class EpochProvingJob implements Traceable {
 
   public async stop(state: EpochProvingJobTerminalState = 'stopped') {
     this.state = state;
-    this.prover.cancel();
+    this.interruptProcessing();
     if (this.runPromise) {
       await this.runPromise;
     }
   }
 
+  private failProcessing() {
+    if (!EpochProvingJobTerminalState.includes(this.state)) {
+      this.state = 'failed';
+    }
+    this.interruptProcessing();
+  }
+
+  private interruptProcessing() {
+    this.abortController.abort();
+    this.prover.cancel();
+  }
+
   private scheduleDeadlineStop() {
     const deadline = this.deadline;
     if (deadline) {
@@ -341,11 +457,14 @@ export class EpochProvingJob implements Traceable {
     const intervalMs = Math.ceil((await l2BlockSource.getL1Constants()).ethereumSlotDuration / 2) * 1000;
     this.epochCheckPromise = new RunningPromise(
       async () => {
-        const blocks = await l2BlockSource.getBlockHeadersForEpoch(this.epochNumber);
-        const blockHashes = await Promise.all(blocks.map(block => block.hash()));
-        const thisBlockHashes = await Promise.all(this.blocks.map(block => block.hash()));
+        const blockHeaders = (
+          await l2BlockSource.getBlocksData({ epoch: this.epochNumber, onlyCheckpointed: true })
+        ).map(d => d.header);
+        const blockHashes = await Promise.all(blockHeaders.map(header => header.hash()));
+        const thisBlocks = this.checkpoints.flatMap(checkpoint => checkpoint.blocks);
+        const thisBlockHashes = await Promise.all(thisBlocks.map(block => block.hash()));
         if (
-          blocks.length !== this.blocks.length ||
+          blockHeaders.length !== thisBlocks.length ||
           !blockHashes.every((block, i) => block.equals(thisBlockHashes[i]))
         ) {
           this.log.warn('Epoch blocks changed underfoot', {
@@ -363,35 +482,27 @@ export class EpochProvingJob implements Traceable {
     this.log.verbose(`Scheduled epoch check for epoch ${this.epochNumber} every ${intervalMs}ms`);
   }
 
-  /* Returns the header for the given block number based on the epoch proving job data. */
-  private getBlockHeader(blockNumber: number) {
-    const block = this.blocks.find(b => b.number === blockNumber);
-    if (block) {
-      return block.getBlockHeader();
-    }
-
-    if (blockNumber === Number(this.data.previousBlockHeader.getBlockNumber())) {
-      return this.data.previousBlockHeader;
-    }
-
-    throw new Error(
-      `Block header not found for block number ${blockNumber} (got ${this.blocks
-        .map(b => b.number)
-        .join(', ')} and previous header ${this.data.previousBlockHeader.getBlockNumber()})`,
-    );
+  /* Returns the last block header in the previous checkpoint for all checkpoints in the epoch */
+  private gatherPreviousBlockHeaders() {
+    const lastBlocks = this.checkpoints.map(checkpoint => checkpoint.blocks.at(-1)!);
+    return [this.data.previousBlockHeader, ...lastBlocks.map(block => block.header).slice(0, -1)];
   }
 
   private getTxs(block: L2Block): Tx[] {
     return block.body.txEffects.map(txEffect => this.txs.get(txEffect.txHash.toString())!);
   }
 
-  private getL1ToL2Messages(block: L2Block) {
-    return this.data.l1ToL2Messages[block.number];
+  private getL1ToL2Messages(checkpoint: Checkpoint) {
+    return this.data.l1ToL2Messages[checkpoint.number];
   }
 
   private async processTxs(publicProcessor: PublicProcessor, txs: Tx[]): Promise<ProcessedTx[]> {
     const { deadline } = this;
-    const [processedTxs, failedTxs] = await publicProcessor.process(txs, { deadline });
+    const [processedTxs, failedTxs] = await publicProcessor.process(txs, {
+      deadline,
+      signal: this.abortController.signal,
+    });
+    this.checkState();
 
     if (failedTxs.length) {
       const failedTxHashes = await Promise.all(failedTxs.map(({ tx }) => tx.getTxHash()));