Integrated tracing into new block proving pipeline

Author: rpanic

packages/sequencer/src/logging/trace.ts

@@ -1,3 +1,4 @@
+// eslint-disable-next-line import/no-extraneous-dependencies
 import snakeCase from "lodash/snakeCase";
 
 import { Tracer } from "./Tracer";
@@ -6,7 +7,7 @@ type Attributes = Record<string, string | number | boolean>;
 
 export function trace<P extends any[], R>(
   name?: string,
-  metadata?: ((p: P /*, r: R*/) => Attributes) | Attributes
+  metadata?: ((p: P) => Attributes) | Attributes
 ) {
   return (
     target: { tracer: Tracer },

packages/sequencer/src/mempool/private/PrivateMempool.ts

@@ -131,8 +131,8 @@ export class PrivateMempool extends SequencerModule implements Mempool {
   // paths are shared between the just succeeded tx and any of the skipped txs. This is
   // because a failed tx may succeed now if the failure was to do with a nonce issue, say.
   // TODO Refactor
-  // eslint-disable-next-line sonarjs/cognitive-complexity
   @trace("mempool.validate_txs")
+  // eslint-disable-next-line sonarjs/cognitive-complexity
   private async checkTxValid(
     transactions: PendingTransaction[],
     baseService: CachedStateService,

packages/sequencer/src/protocol/production/BatchProducerModule.ts

@@ -17,12 +17,10 @@ import { CachedMerkleTreeStore } from "../../state/merkle/CachedMerkleTreeStore"
 import { AsyncMerkleTreeStore } from "../../state/async/AsyncMerkleTreeStore";
 import { BlockWithResult } from "../../storage/model/Block";
 import type { Database } from "../../storage/Database";
-import { Tracer } from "../../logging/Tracer";
 
 import { BlockProofSerializer } from "./tasks/serializers/BlockProofSerializer";
 import { BatchTracingService } from "./tracing/BatchTracingService";
 import { BatchFlow } from "./flow/BatchFlow";
-import { trace } from "../../logging/trace";
 
 export type StateRecord = Record<string, Field[] | undefined>;
 
@@ -54,8 +52,6 @@ export class BatchProducerModule extends SequencerModule {
     @inject("BatchStorage") private readonly batchStorage: BatchStorage,
     @inject("Database")
     private readonly database: Database,
-    @inject("Tracer")
-    public readonly tracer: Tracer,
     private readonly batchFlow: BatchFlow,
     private readonly blockProofSerializer: BlockProofSerializer,
     private readonly batchTraceService: BatchTracingService
@@ -139,7 +135,6 @@ export class BatchProducerModule extends SequencerModule {
     noop();
   }
 
-  @trace("batch")
   private async produceBatch(
     blocks: BlockWithResult[],
     height: number
@@ -175,12 +170,12 @@ export class BatchProducerModule extends SequencerModule {
    *
    *
    * @param blocks
-   * @param blockId
+   * @param batchId
    * @private
    */
   private async computeBatch(
     blocks: BlockWithResult[],
-    blockId: number
+    batchId: number
   ): Promise<{
     proof: Proof<BlockProverPublicInput, BlockProverPublicOutput>;
     changes: CachedMerkleTreeStore;
@@ -195,10 +190,11 @@ export class BatchProducerModule extends SequencerModule {
 
     const trace = await this.batchTraceService.traceBatch(
       blocks.map((block) => block),
-      merkleTreeStore
+      merkleTreeStore,
+      batchId
     );
 
-    const proof = await this.batchFlow.executeBatch(trace, blockId);
+    const proof = await this.batchFlow.executeBatch(trace, batchId);
 
     const fromNetworkState = blocks[0].block.networkState.before;
     const toNetworkState = blocks.at(-1)!.result.afterNetworkState;

packages/sequencer/src/protocol/production/flow/BatchFlow.ts

@@ -18,6 +18,8 @@ import { FlowCreator } from "../../../worker/flow/Flow";
 import { NewBlockProvingParameters, NewBlockTask } from "../tasks/NewBlockTask";
 import { BlockReductionTask } from "../tasks/BlockReductionTask";
 import { BatchTrace } from "../tracing/BatchTracingService";
+import { Tracer } from "../../../logging/Tracer";
+import { trace } from "../../../logging/trace";
 
 import { ReductionTaskFlow } from "./ReductionTaskFlow";
 import { StateTransitionFlow } from "./StateTransitionFlow";
@@ -33,7 +35,9 @@ export class BatchFlow {
     private readonly stateTransitionFlow: StateTransitionFlow,
     private readonly blockFlow: BlockFlow,
     @inject("Protocol")
-    private readonly protocol: Protocol<MandatoryProtocolModulesRecord>
+    private readonly protocol: Protocol<MandatoryProtocolModulesRecord>,
+    @inject("Tracer")
+    public readonly tracer: Tracer
   ) {}
 
   private isBlockProofsMergable(a: BlockProof, b: BlockProof): boolean {
@@ -88,6 +92,7 @@ export class BatchFlow {
     );
   }
 
+  @trace("batch.prove", ([, batchId]) => ({ batchId }))
   public async executeBatch(batch: BatchTrace, batchId: number) {
     const batchFlow = new ReductionTaskFlow(
       {
@@ -104,9 +109,13 @@ export class BatchFlow {
       number,
       Nullable<NewBlockProvingParameters>
     > = Object.fromEntries(
-      batch.blocks.map((trace, i) => [
+      batch.blocks.map((blockTrace, i) => [
         i,
-        { params: trace.blockParams, input1: undefined, input2: undefined },
+        {
+          params: blockTrace.blockParams,
+          input1: undefined,
+          input2: undefined,
+        },
       ])
     );
 
@@ -125,8 +134,8 @@ export class BatchFlow {
       }
     );
 
-    await mapSequential(batch.blocks, async (trace, blockIndex) => {
-      await this.blockFlow.executeBlock(trace, async (proof) => {
+    await mapSequential(batch.blocks, async (blockTrace, blockIndex) => {
+      await this.blockFlow.executeBlock(blockTrace, async (proof) => {
         map[blockIndex].input2 = proof;
         await this.pushBlockInput(map[blockIndex], batchFlow);
       });

packages/sequencer/src/protocol/production/sequencing/BlockProducerModule.ts

@@ -26,10 +26,10 @@ import {
 import { MessageStorage } from "../../../storage/repositories/MessageStorage";
 import { Database } from "../../../storage/Database";
 import { Tracer } from "../../../logging/Tracer";
+import { trace } from "../../../logging/trace";
 
 import { BlockProductionService } from "./BlockProductionService";
 import { BlockResultService } from "./BlockResultService";
-import { trace } from "../../../logging/trace";
 
 export interface BlockConfig {
   allowEmptyBlock?: boolean;
@@ -109,7 +109,7 @@ export class BlockProducerModule extends SequencerModule<BlockConfig> {
     }
   }
 
-  @trace("block.metadata", ([block]) => ({ height: block.height.toString() }))
+  @trace("block.result", ([block]) => ({ height: block.height.toString() }))
   public async generateMetadata(block: Block): Promise<BlockResult> {
     const traceMetadata = {
       height: block.height.toString(),
@@ -123,13 +123,18 @@ export class BlockProducerModule extends SequencerModule<BlockConfig> {
         this.unprovenStateService
       );
 
-    await this.database.executeInTransaction(async () => {
-      await blockHashTreeStore.mergeIntoParent();
-      await treeStore.mergeIntoParent();
-      await stateService.mergeIntoParent();
-
-      await this.blockQueue.pushResult(result);
-    });
+    await this.tracer.trace(
+      "block.result.commit",
+      async () =>
+        await this.database.executeInTransaction(async () => {
+          await blockHashTreeStore.mergeIntoParent();
+          await treeStore.mergeIntoParent();
+          await stateService.mergeIntoParent();
+
+          await this.blockQueue.pushResult(result);
+        }),
+      traceMetadata
+    );
 
     return result;
   }

packages/sequencer/src/protocol/production/sequencing/BlockProductionService.ts

@@ -25,13 +25,14 @@ import { CachedStateService } from "../../../state/state/CachedStateService";
 import { PendingTransaction } from "../../../mempool/PendingTransaction";
 import { AsyncStateService } from "../../../state/async/AsyncStateService";
 import { UntypedStateTransition } from "../helpers/UntypedStateTransition";
+import { Tracer } from "../../../logging/Tracer";
+import { trace } from "../../../logging/trace";
 
 import {
   BlockTrackers,
   executeWithExecutionContext,
   TransactionExecutionService,
 } from "./TransactionExecutionService";
-import { Tracer } from "../../../logging/Tracer";
 
 @injectable()
 @scoped(Lifecycle.ContainerScoped)
@@ -42,7 +43,7 @@ export class BlockProductionService {
     @inject("Protocol")
     protocol: Protocol<MandatoryProtocolModulesRecord & ProtocolModulesRecord>,
     @inject("Tracer")
-    private readonly tracer: Tracer,
+    public readonly tracer: Tracer,
     private readonly transactionExecutionService: TransactionExecutionService,
     @inject("StateServiceProvider")
     private readonly stateServiceProvider: StateServiceProvider
@@ -51,6 +52,7 @@ export class BlockProductionService {
       protocol.dependencyContainer.resolveAll("ProvableBlockHook");
   }
 
+  @trace("block.hook.before")
   public async executeBeforeBlockHook(
     args: BeforeBlockHookArguments,
     inputNetworkState: NetworkState,

packages/sequencer/src/protocol/production/sequencing/BlockResultService.ts

@@ -25,6 +25,8 @@ import { UntypedStateTransition } from "../helpers/UntypedStateTransition";
 import { CachedStateService } from "../../../state/state/CachedStateService";
 import { AsyncStateService } from "../../../state/async/AsyncStateService";
 import type { StateRecord } from "../BatchProducerModule";
+import { trace } from "../../../logging/trace";
+import { Tracer } from "../../../logging/Tracer";
 
 import { executeWithExecutionContext } from "./TransactionExecutionService";
 
@@ -72,12 +74,15 @@ export class BlockResultService {
     @inject("Protocol")
     protocol: Protocol<MandatoryProtocolModulesRecord & ProtocolModulesRecord>,
     @inject("StateServiceProvider")
-    private readonly stateServiceProvider: StateServiceProvider
+    private readonly stateServiceProvider: StateServiceProvider,
+    @inject("Tracer")
+    public readonly tracer: Tracer
   ) {
     this.blockHooks =
       protocol.dependencyContainer.resolveAll("ProvableBlockHook");
   }
 
+  @trace("block.hook.after")
   public async executeAfterBlockHook(
     args: AfterBlockHookArguments,
     inputNetworkState: NetworkState,
@@ -169,6 +174,9 @@ export class BlockResultService {
     return tree;
   }
 
+  @trace("block.result.generate", ([block]) => ({
+    height: block.height.toString(),
+  }))
   public async generateMetadataForNextBlock(
     block: Block,
     merkleTreeStore: AsyncMerkleTreeStore,

packages/sequencer/src/protocol/production/tracing/BatchTracingService.ts

@@ -6,11 +6,13 @@ import {
   WitnessedRootHashList,
 } from "@proto-kit/protocol";
 import { Field } from "o1js";
-import { injectable } from "tsyringe";
+import { inject, injectable } from "tsyringe";
 
 import { CachedMerkleTreeStore } from "../../../state/merkle/CachedMerkleTreeStore";
 import { StateTransitionProofParameters } from "../tasks/StateTransitionTask";
 import { BlockWithResult } from "../../../storage/model/Block";
+import { trace } from "../../../logging/trace";
+import { Tracer } from "../../../logging/Tracer";
 
 import {
   BlockTrace,
@@ -30,7 +32,9 @@ export type BatchTrace = {
 export class BatchTracingService {
   public constructor(
     private readonly blockTracingService: BlockTracingService,
-    private readonly stateTransitionTracingService: StateTransitionTracingService
+    private readonly stateTransitionTracingService: StateTransitionTracingService,
+    @inject("Tracer")
+    public readonly tracer: Tracer
   ) {}
 
   private createBatchState(block: BlockWithResult): BatchTracingState {
@@ -46,6 +50,7 @@ export class BatchTracingService {
     };
   }
 
+  @trace("batch.trace.blocks")
   public async traceBlocks(blocks: BlockWithResult[]) {
     const batchState = this.createBatchState(blocks[0]);
 
@@ -58,34 +63,42 @@ export class BatchTracingService {
           ...state,
           transactionList: new TransactionHashList(),
         };
-        const [newState, trace] = await this.blockTracingService.traceBlock(
-          blockProverState,
-          block,
-          index === numBlocks - 1
-        );
-        return [newState, trace];
+        const [newState, blockTrace] =
+          await this.blockTracingService.traceBlock(
+            blockProverState,
+            block,
+            index === numBlocks - 1
+          );
+        return [newState, blockTrace];
       },
       batchState
     );
 
     return blockTraces;
   }
 
+  @trace("batch.trace.transitions")
   public async traceStateTransitions(
     blocks: BlockWithResult[],
     merkleTreeStore: CachedMerkleTreeStore
   ) {
-    const batches = this.stateTransitionTracingService.extractSTBatches(blocks);
+    const batches = await this.tracer.trace(
+      "batch.trace.transitions.encoding",
+      async () => this.stateTransitionTracingService.extractSTBatches(blocks)
+    );
 
     return await this.stateTransitionTracingService.createMerkleTrace(
       merkleTreeStore,
       batches
     );
   }
 
+  @trace("batch.trace", ([, , batchId]) => ({ batchId }))
   public async traceBatch(
     blocks: BlockWithResult[],
-    merkleTreeStore: CachedMerkleTreeStore
+    merkleTreeStore: CachedMerkleTreeStore,
+    // Only for trace metadata
+    batchId: number
   ): Promise<BatchTrace> {
     if (blocks.length === 0) {
       return { blocks: [], stateTransitionTrace: [] };

packages/sequencer/src/protocol/production/tracing/BlockTracingService.ts

@@ -8,10 +8,12 @@ import { toStateTransitionsHash } from "@proto-kit/module";
 import { yieldSequential } from "@proto-kit/common";
 // eslint-disable-next-line import/no-extraneous-dependencies
 import chunk from "lodash/chunk";
-import { injectable } from "tsyringe";
+import { inject, injectable } from "tsyringe";
 
 import { BlockWithResult } from "../../../storage/model/Block";
 import type { NewBlockProverParameters } from "../tasks/NewBlockTask";
+import { Tracer } from "../../../logging/Tracer";
+import { trace } from "../../../logging/trace";
 
 import {
   collectStartingState,
@@ -42,9 +44,14 @@ export type BlockTrace = {
 @injectable()
 export class BlockTracingService {
   public constructor(
-    private readonly transactionTracing: TransactionTracingService
+    private readonly transactionTracing: TransactionTracingService,
+    @inject("Tracer")
+    public readonly tracer: Tracer
   ) {}
 
+  @trace("batch.trace.block", ([, block]) => ({
+    height: block.block.height.toString(),
+  }))
   public async traceBlock(
     state: BlockTracingState,
     block: BlockWithResult,
@@ -85,7 +92,7 @@ export class BlockTracingService {
     const [afterState, transactionTraces] = await yieldSequential(
       chunk(block.block.transactions, 2),
       async (input, [transaction1, transaction2]) => {
-        const [output, trace] =
+        const [output, transactionTrace] =
           transaction2 !== undefined
             ? await this.transactionTracing.createMultiTransactionTrace(
                 input,
@@ -97,7 +104,7 @@ export class BlockTracingService {
                 transaction1
               );
 
-        return [output, trace];
+        return [output, transactionTrace];
       },
       state
     );