refactor: moving BenchmarkedNode to e2e

yarn-project/cli-wallet/src/utils/profiling.ts

@@ -93,41 +93,6 @@ export function printProfileResult(
     );
   }
 
-  if (stats.nodeRPCCalls) {
-    log(format('\nRPC calls:\n'));
-    for (const [method, { times }] of Object.entries(stats.nodeRPCCalls.perMethod)) {
-      const calls = times.length;
-      const total = times.reduce((acc, time) => acc + time, 0);
-      const avg = total / calls;
-      const min = Math.min(...times);
-      const max = Math.max(...times);
-      log(
-        format(
-          method.padEnd(ORACLE_NAME_PADDING),
-          `${calls} calls`.padStart(COLUMN_MIN_WIDTH).padEnd(COLUMN_MAX_WIDTH),
-          `${total.toFixed(2)}ms`.padStart(COLUMN_MIN_WIDTH).padEnd(COLUMN_MAX_WIDTH),
-          `min: ${min.toFixed(2)}ms`.padStart(COLUMN_MIN_WIDTH).padEnd(COLUMN_MAX_WIDTH),
-          `avg: ${avg.toFixed(2)}ms`.padStart(COLUMN_MIN_WIDTH).padEnd(COLUMN_MAX_WIDTH),
-          `max: ${max.toFixed(2)}ms`.padStart(COLUMN_MIN_WIDTH).padEnd(COLUMN_MAX_WIDTH),
-        ),
-      );
-    }
-
-    const { roundTrips } = stats.nodeRPCCalls;
-    log(format('\nRound trips (actual blocking waits):\n'));
-    log(format('Round trips:'.padEnd(25), `${roundTrips.roundTrips}`.padStart(COLUMN_MAX_WIDTH)));
-    log(
-      format(
-        'Total blocking time:'.padEnd(25),
-        `${roundTrips.totalBlockingTime.toFixed(2)}ms`.padStart(COLUMN_MAX_WIDTH),
-      ),
-    );
-    if (roundTrips.roundTrips > 0) {
-      const avgRoundTrip = roundTrips.totalBlockingTime / roundTrips.roundTrips;
-      log(format('Avg round trip:'.padEnd(25), `${avgRoundTrip.toFixed(2)}ms`.padStart(COLUMN_MAX_WIDTH)));
-    }
-  }
-
   log(format('\nSync time:'.padEnd(25), `${timings.sync?.toFixed(2)}ms`.padStart(16)));
   log(
     format(

yarn-project/end-to-end/src/bench/client_flows/amm.test.ts

@@ -22,20 +22,20 @@ const AMOUNT_PER_NOTE = 1_000_000;
 const MINIMUM_NOTES_FOR_RECURSION_LEVEL = [0, 2, 10];
 
 // Set to true to print out the round trip information to the console.
-const DEBUG_ROUND_TRIPS = true;
+const DEBUG_ROUND_TRIPS = false;
 
 // Expected number of node round trips per account contract and payment method.
 const EXPECTED_ROUND_TRIPS: Record<string, number> = {
-  'ecdsar1+private_fpc': 118,
-  'ecdsar1+sponsored_fpc': 80,
-  'schnorr+private_fpc': 118,
-  'schnorr+sponsored_fpc': 82,
+  'ecdsar1+private_fpc': 165,
+  'ecdsar1+sponsored_fpc': 12,
+  'schnorr+private_fpc': 143,
+  'schnorr+sponsored_fpc': 217,
 };
 
 interface RoundTripData {
   accountType: AccountType;
   paymentMethod: BenchmarkingFeePaymentMethod;
-  roundTrips: RoundTripStats | undefined;
+  roundTrips: RoundTripStats;
 }
 
 describe('AMM benchmark', () => {
@@ -180,7 +180,10 @@ describe('AMM benchmark', () => {
               .methods.add_liquidity(amountToSend, amountToSend, amountToSend, amountToSend, nonceForAuthwits)
               .with({ authWitnesses: [token0Authwit, token1Authwit] });
 
-            const profileResult = await captureProfile(
+            // Get node stats from the benchmarked node to track RPC calls
+            const nodeStats = t.benchmarkedNode.getAndResetStats();
+
+            await captureProfile(
               `${accountType}+amm_add_liquidity_1_recursions+${benchmarkingPaymentMethod}`,
               addLiquidityInteraction,
               options,
@@ -196,6 +199,7 @@ describe('AMM benchmark', () => {
                 1 + // Kernel reset
                 1 + // Kernel tail
                 1, // Kernel hiding
+              nodeStats,
             );
 
             if (process.env.SANITY_CHECKS) {
@@ -207,11 +211,11 @@ describe('AMM benchmark', () => {
               roundTripData.push({
                 accountType,
                 paymentMethod: benchmarkingPaymentMethod,
-                roundTrips: profileResult.stats.nodeRPCCalls?.roundTrips,
+                roundTrips: nodeStats.roundTrips,
               });
             } else {
               const roundTripsKey = `${accountType}+${benchmarkingPaymentMethod}`;
-              const actualRoundTrips = profileResult.stats.nodeRPCCalls?.roundTrips.roundTrips ?? 0;
+              const actualRoundTrips = nodeStats.roundTrips.roundTrips;
               const expectedRoundTrips = EXPECTED_ROUND_TRIPS[roundTripsKey];
               if (expectedRoundTrips === undefined) {
                 throw new Error(
@@ -246,14 +250,6 @@ describe('AMM benchmark', () => {
    * Prints out the round trip information that can be used to debug unexpected round trip changes.
    */
   function printRoundTripDebuggingInfo(data: RoundTripData[]) {
-    for (const trip of data) {
-      if (!trip.roundTrips) {
-        throw new Error(
-          `Round trip stats are undefined for ${trip.accountType}+${trip.paymentMethod}. This should not happen.`,
-        );
-      }
-    }
-
     const width = 120; // Fixed standard width
     const title = '  ROUND TRIP DEBUGGING INFORMATION';
 
@@ -325,13 +321,7 @@ describe('AMM benchmark', () => {
 
       for (let itemIdx = 0; itemIdx < items.length; itemIdx++) {
         const item = items[itemIdx];
-
-        if (!item.roundTrips) {
-          throw new Error(
-            `Round trip stats are undefined for ${item.accountType}+${item.paymentMethod}. This should not happen.`,
-          );
-        }
-        const roundTrips = item.roundTrips; // TypeScript now knows it's defined
+        const roundTrips = item.roundTrips;
 
         if (itemIdx > 0) {
           output += leftBorder + ' '.repeat(width - 2) + rightBorder + '\n';

yarn-project/end-to-end/src/bench/client_flows/benchmark.ts

@@ -8,6 +8,7 @@ import type { Logger } from '@aztec/aztec.js/log';
 import { createLogger } from '@aztec/foundation/log';
 import { type PrivateExecutionStep, serializePrivateExecutionSteps } from '@aztec/stdlib/kernel';
 import type {
+  NodeStats,
   ProvingStats,
   ProvingTimings,
   RoundTripStats,
@@ -180,6 +181,7 @@ export function generateBenchmark(
   privateExecutionSteps: PrivateExecutionStep[],
   proverType: ProverType,
   error: string | undefined,
+  nodeStats?: NodeStats,
 ): ClientFlowBenchmark {
   let maxMemory = 0;
   let minimumTrace: StructuredTrace;
@@ -219,7 +221,7 @@ export function generateBenchmark(
   }, []);
   const timings = stats.timings;
   const totalGateCount = steps[steps.length - 1].accGateCount;
-  const nodeRPCCalls = stats.nodeRPCCalls ?? {
+  const nodeRPCCalls = nodeStats ?? {
     perMethod: {},
     roundTrips: { roundTrips: 0, totalBlockingTime: 0, roundTripDurations: [], roundTripMethods: [] },
   };
@@ -325,6 +327,7 @@ export async function captureProfile(
   interaction: ContractFunctionInteraction | DeployMethod,
   opts: Omit<ProfileInteractionOptions & DeployOptions, 'profileMode'>,
   expectedSteps?: number,
+  nodeStats?: NodeStats,
 ) {
   // Make sure the proxy logger starts from a clean slate
   ProxyLogger.getInstance().flushLogs();
@@ -333,7 +336,7 @@ export async function captureProfile(
   if (expectedSteps !== undefined && result.executionSteps.length !== expectedSteps) {
     throw new Error(`Expected ${expectedSteps} execution steps, got ${result.executionSteps.length}`);
   }
-  const benchmark = generateBenchmark(label, logs, result.stats, result.executionSteps, 'wasm', undefined);
+  const benchmark = generateBenchmark(label, logs, result.stats, result.executionSteps, 'wasm', undefined, nodeStats);
 
   const ivcFolder = process.env.CAPTURE_IVC_FOLDER;
   if (ivcFolder) {

yarn-project/pxe/src/contract_function_simulator/benchmarked_node.ts → yarn-project/end-to-end/src/bench/client_flows/benchmarked_node.ts

@@ -17,7 +17,7 @@ import type { NodeStats, RoundTripStats } from '@aztec/stdlib/tx';
  * value in `safe_json_rpc_client.ts` to 1 (the main motivation for batching was to get around parallel http requests
  * limits in web browsers which is not a problem when debugging in node.js).
  */
-export type BenchmarkedNode = AztecNode & { getStats(): NodeStats };
+export type BenchmarkedNode = AztecNode & { getAndResetStats(): NodeStats };
 
 export class BenchmarkedNodeFactory {
   static create(node: AztecNode): BenchmarkedNode {
@@ -37,9 +37,19 @@ export class BenchmarkedNodeFactory {
 
     return new Proxy(node, {
       get(target, prop: keyof BenchmarkedNode) {
-        if (prop === 'getStats') {
+        if (prop === 'getAndResetStats') {
           return (): NodeStats => {
-            return { perMethod, roundTrips };
+            const stats = { perMethod: { ...perMethod }, roundTrips: { ...roundTrips } };
+            // Reset perMethod
+            for (const key of Object.keys(perMethod)) {
+              delete perMethod[key as keyof AztecNode];
+            }
+            // Reset roundTrips
+            roundTrips.roundTrips = 0;
+            roundTrips.totalBlockingTime = 0;
+            roundTrips.roundTripDurations = [];
+            roundTrips.roundTripMethods = [];
+            return stats;
           };
         } else {
           return function (...args: any[]) {

yarn-project/end-to-end/src/bench/client_flows/client_flows_benchmark.ts

@@ -38,6 +38,7 @@ import {
   type GasBridgingTestHarness,
 } from '../../shared/gas_portal_test_harness.js';
 import { ProxyLogger } from './benchmark.js';
+import { type BenchmarkedNode, BenchmarkedNodeFactory } from './benchmarked_node.js';
 import { type ClientFlowsConfig, FULL_FLOWS_CONFIG, KEY_FLOWS_CONFIG } from './config.js';
 
 const { BENCHMARK_CONFIG } = process.env;
@@ -85,6 +86,9 @@ export class ClientFlowsBenchmark {
   // PXE and Wallet used by the benchmarking user. It can be set up with client-side proving enabled
   public userWallet!: TestWallet;
 
+  // Benchmarked node that tracks RPC call stats (wraps aztecNode)
+  public benchmarkedNode!: BenchmarkedNode;
+
   public realProofs = ['true', '1'].includes(process.env.REAL_PROOFS ?? '');
 
   public paymentMethods: Record<BenchmarkingFeePaymentMethod, { forWallet: FeePaymentMethodGetter; circuits: number }> =
@@ -224,7 +228,10 @@ export class ClientFlowsBenchmark {
       proverEnabled: this.realProofs,
     } as PXEConfig;
 
-    this.userWallet = await TestWallet.create(this.aztecNode, userPXEConfigWithContracts, {
+    // Wrap the node with BenchmarkedNode to track RPC call stats
+    this.benchmarkedNode = BenchmarkedNodeFactory.create(this.aztecNode);
+
+    this.userWallet = await TestWallet.create(this.benchmarkedNode, userPXEConfigWithContracts, {
       loggers: {
         prover: this.proxyLogger.createLogger('pxe:bb:wasm:bundle:proxied'),
       },

yarn-project/pxe/src/contract_function_simulator/contract_function_simulator.ts

@@ -80,7 +80,6 @@ import type { PrivateEventStore } from '../storage/private_event_store/private_e
 import type { RecipientTaggingStore } from '../storage/tagging_store/recipient_tagging_store.js';
 import type { SenderAddressBookStore } from '../storage/tagging_store/sender_address_book_store.js';
 import type { SenderTaggingStore } from '../storage/tagging_store/sender_tagging_store.js';
-import type { BenchmarkedNode } from './benchmarked_node.js';
 import { ExecutionNoteCache } from './execution_note_cache.js';
 import { ExecutionTaggingIndexCache } from './execution_tagging_index_cache.js';
 import { HashedValuesCache } from './hashed_values_cache.js';
@@ -325,22 +324,6 @@ export class ContractFunctionSimulator {
     }
   }
   // docs:end:execute_utility_function
-
-  /**
-   * Returns the execution statistics collected during the simulator run.
-   * @returns The execution statistics.
-   */
-  getStats() {
-    const nodeRPCCalls =
-      typeof (this.aztecNode as BenchmarkedNode).getStats === 'function'
-        ? (this.aztecNode as BenchmarkedNode).getStats()
-        : {
-            perMethod: {},
-            roundTrips: { roundTrips: 0, totalBlockingTime: 0, roundTripDurations: [], roundTripMethods: [] },
-          };
-
-    return { nodeRPCCalls };
-  }
 }
 
 class OrderedSideEffect<T> {

yarn-project/pxe/src/pxe.ts

@@ -53,7 +53,6 @@ import { inspect } from 'util';
 
 import { BlockSynchronizer } from './block_synchronizer/index.js';
 import type { PXEConfig } from './config/index.js';
-import { BenchmarkedNodeFactory } from './contract_function_simulator/benchmarked_node.js';
 import {
   ContractFunctionSimulator,
   generateSimulatedProvingResult,
@@ -206,7 +205,7 @@ export class PXE {
       this.noteStore,
       this.keyStore,
       this.addressStore,
-      BenchmarkedNodeFactory.create(this.node),
+      this.node,
       this.anchorBlockStore,
       this.senderTaggingStore,
       this.recipientTaggingStore,
@@ -660,7 +659,6 @@ export class PXE {
 
         const txProvingResult = new TxProvingResult(privateExecutionResult, publicInputs, chonkProof!, {
           timings,
-          nodeRPCCalls: contractFunctionSimulator?.getStats().nodeRPCCalls,
         });
 
         // While not strictly necessary to store tagging cache contents in the DB since we sync tagging indexes from
@@ -766,8 +764,7 @@ export class PXE {
             total - ((syncTime ?? 0) + (proving ?? 0) + perFunction.reduce((acc, { time }) => acc + time, 0)),
         };
 
-        const simulatorStats = contractFunctionSimulator.getStats();
-        return new TxProfileResult(executionSteps, { timings, nodeRPCCalls: simulatorStats.nodeRPCCalls });
+        return new TxProfileResult(executionSteps, { timings });
       } catch (err: any) {
         throw this.#contextualizeError(err, inspect(txRequest), `profileMode=${profileMode}`);
       }
@@ -909,10 +906,8 @@ export class PXE {
             : {}),
         });
 
-        const simulatorStats = contractFunctionSimulator.getStats();
         return TxSimulationResult.fromPrivateSimulationResultAndPublicOutput(privateSimulationResult, publicOutput, {
           timings,
-          nodeRPCCalls: simulatorStats.nodeRPCCalls,
         });
       } catch (err: any) {
         throw this.#contextualizeError(
@@ -976,8 +971,7 @@ export class PXE {
           unaccounted: totalTime - (syncTime + perFunction.reduce((acc, { time }) => acc + time, 0)),
         };
 
-        const simulationStats = contractFunctionSimulator.getStats();
-        return { result: executionResult, stats: { timings, nodeRPCCalls: simulationStats.nodeRPCCalls } };
+        return { result: executionResult, stats: { timings } };
       } catch (err: any) {
         const { to, name, args } = call;
         const stringifiedArgs = args.map(arg => arg.toString()).join(', ');