refactor(tinybench): share structure across walltime and instrumented

Author: GuillaumeLagrange

packages/tinybench-plugin/src/index.ts

@@ -3,7 +3,6 @@ import {
   getGitDir,
   InstrumentHooks,
   mongoMeasurement,
-  setupCore,
   SetupInstrumentsRequestBody,
   SetupInstrumentsResponse,
   tryIntrospect,
@@ -12,9 +11,9 @@ import path from "path";
 import { get as getStackTrace } from "stack-trace";
 import { Bench } from "tinybench";
 import { fileURLToPath } from "url";
-import { runInstrumentedBench } from "./instrumented";
+import { setupCodspeedInstrumentedBench } from "./instrumented";
 import { getOrCreateUriMap } from "./uri";
-import { runWalltimeBench } from "./walltime";
+import { setupCodspeedWalltimeBench } from "./walltime";
 
 tryIntrospect();
 
@@ -23,7 +22,6 @@ export function withCodSpeed(bench: Bench): Bench {
   if (codspeedRunnerMode === "disabled") {
     return bench;
   }
-  setupCore();
 
   const rootCallingFile = getCallingFile();
 
@@ -42,9 +40,9 @@ export function withCodSpeed(bench: Bench): Bench {
   };
 
   if (codspeedRunnerMode === "instrumented") {
-    runInstrumentedBench(bench, rootCallingFile);
+    setupCodspeedInstrumentedBench(bench, rootCallingFile);
   } else if (codspeedRunnerMode === "walltime") {
-    runWalltimeBench(bench, rootCallingFile);
+    setupCodspeedWalltimeBench(bench, rootCallingFile);
   }
 
   return bench;

packages/tinybench-plugin/src/index.unit.test.ts

@@ -3,15 +3,15 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { withCodSpeed } from ".";
 
 const mockInstrumented = vi.hoisted(() => ({
-  runInstrumentedBench: vi.fn(),
+  setupCodspeedInstrumentedBench: vi.fn(),
 }));
 
 vi.mock("./instrumented", () => ({
   ...mockInstrumented,
 }));
 
 const mockWalltime = vi.hoisted(() => ({
-  runWalltimeBench: vi.fn(),
+  setupCodspeedWalltimeBench: vi.fn(),
 }));
 
 vi.mock("./walltime", () => ({
@@ -44,8 +44,8 @@ describe("withCodSpeed behavior without different codspeed modes", () => {
 
     withCodSpeed(new Bench());
 
-    expect(mockInstrumented.runInstrumentedBench).toHaveBeenCalled();
-    expect(mockWalltime.runWalltimeBench).not.toHaveBeenCalled();
+    expect(mockInstrumented.setupCodspeedInstrumentedBench).toHaveBeenCalled();
+    expect(mockWalltime.setupCodspeedWalltimeBench).not.toHaveBeenCalled();
   });
 
   it("should run in walltime mode when CODSPEED_RUNNER_MODE=walltime", async () => {
@@ -54,7 +54,9 @@ describe("withCodSpeed behavior without different codspeed modes", () => {
 
     withCodSpeed(new Bench());
 
-    expect(mockInstrumented.runInstrumentedBench).not.toHaveBeenCalled();
-    expect(mockWalltime.runWalltimeBench).toHaveBeenCalled();
+    expect(
+      mockInstrumented.setupCodspeedInstrumentedBench
+    ).not.toHaveBeenCalled();
+    expect(mockWalltime.setupCodspeedWalltimeBench).toHaveBeenCalled();
   });
 });

packages/tinybench-plugin/src/instrumented.ts

@@ -2,18 +2,40 @@ import {
   InstrumentHooks,
   mongoMeasurement,
   optimizeFunction,
-  teardownCore,
 } from "@codspeed/core";
 import { Bench, Fn, FnOptions, Task } from "tinybench";
-import { getTaskUri } from "./uri";
+import { BaseBenchRunner } from "./shared";
 
-declare const __VERSION__: string;
-
-export function runInstrumentedBench(
+export function setupCodspeedInstrumentedBench(
   bench: Bench,
   rootCallingFile: string
 ): void {
-  const runTaskAsync = async (task: Task, uri: string): Promise<void> => {
+  const runner = new InstrumentedBenchRunner(bench, rootCallingFile);
+  runner.setupBenchMethods();
+}
+
+class InstrumentedBenchRunner extends BaseBenchRunner {
+  protected getModeName(): string {
+    return "instrumented mode";
+  }
+
+  private taskCompletionMessage() {
+    return InstrumentHooks.isInstrumented() ? "Measured" : "Checked";
+  }
+
+  private wrapFunctionWithFrame(fn: Fn, isAsync: boolean): Fn {
+    if (isAsync) {
+      return async function __codspeed_root_frame__() {
+        await fn();
+      };
+    } else {
+      return function __codspeed_root_frame__() {
+        fn();
+      };
+    }
+  }
+
+  protected async runTaskAsync(task: Task, uri: string): Promise<void> {
     const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
     await fnOpts?.beforeAll?.call(task, "run");
@@ -25,79 +47,38 @@ export function runInstrumentedBench(
     await fnOpts?.beforeEach?.call(task, "run");
     await mongoMeasurement.start(uri);
 
-    await (async function __codspeed_root_frame__() {
-      global.gc?.();
-      InstrumentHooks.startBenchmark();
-      await fn();
-      InstrumentHooks.stopBenchmark();
-      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-    })();
+    global.gc?.();
+    await this.wrapWithInstrumentHooksAsync(
+      this.wrapFunctionWithFrame(fn, true),
+      uri
+    );
 
     await mongoMeasurement.stop(uri);
     await fnOpts?.afterEach?.call(task, "run");
     await fnOpts?.afterAll?.call(task, "run");
-  };
 
-  // Sync task runner
-  const runTaskSync = (task: Task, uri: string): void => {
+    this.logTaskCompletion(uri, this.taskCompletionMessage());
+  }
+
+  protected runTaskSync(task: Task, uri: string): void {
     const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
     fnOpts?.beforeAll?.call(task, "run");
     fnOpts?.beforeEach?.call(task, "run");
 
-    (function __codspeed_root_frame__() {
-      global.gc?.();
-      InstrumentHooks.startBenchmark();
-      fn();
-      InstrumentHooks.stopBenchmark();
-      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-    })();
+    this.wrapWithInstrumentHooks(this.wrapFunctionWithFrame(fn, false), uri);
 
     fnOpts?.afterEach?.call(task, "run");
     fnOpts?.afterAll?.call(task, "run");
-  };
-
-  bench.run = async () => {
-    logStart();
 
-    for (const task of bench.tasks) {
-      const uri = getTaskUri(bench, task.name, rootCallingFile);
-      await runTaskAsync(task, uri);
-      logTaskCompletion(uri);
-    }
-
-    return logEnd();
-  };
-
-  bench.runSync = () => {
-    logStart();
-
-    for (const task of bench.tasks) {
-      const uri = getTaskUri(bench, task.name, rootCallingFile);
-      runTaskSync(task, uri);
-      logTaskCompletion(uri);
-    }
+    this.logTaskCompletion(uri, this.taskCompletionMessage());
+  }
 
-    return logEnd();
-  };
-
-  const logStart = () => {
-    console.log(
-      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
-    );
-  };
-
-  const logTaskCompletion = (uri: string) => {
-    console.log(
-      `    ✔ ${
-        InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
-      } ${uri}`
-    );
-  };
+  protected finalizeAsyncRun(): Task[] {
+    return this.finalizeBenchRun();
+  }
 
-  const logEnd = () => {
-    teardownCore();
-    console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);
-    return bench.tasks;
-  };
+  protected finalizeSyncRun(): Task[] {
+    return this.finalizeBenchRun();
+  }
 }

packages/tinybench-plugin/src/shared.ts

@@ -0,0 +1,89 @@
+import { InstrumentHooks, setupCore, teardownCore } from "@codspeed/core";
+import { Bench, Fn, Task } from "tinybench";
+import { getTaskUri } from "./uri";
+
+declare const __VERSION__: string;
+
+export abstract class BaseBenchRunner {
+  protected bench: Bench;
+  protected rootCallingFile: string;
+
+  constructor(bench: Bench, rootCallingFile: string) {
+    this.bench = bench;
+    this.rootCallingFile = rootCallingFile;
+  }
+
+  private setupBenchRun(): void {
+    setupCore();
+    this.logStart();
+  }
+
+  private logStart(): void {
+    console.log(
+      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (${this.getModeName()})`
+    );
+  }
+
+  protected getTaskUri(task: Task): string {
+    return getTaskUri(this.bench, task.name, this.rootCallingFile);
+  }
+
+  protected logTaskCompletion(uri: string, status: string): void {
+    console.log(`[CodSpeed] ${status} ${uri}`);
+  }
+
+  protected finalizeBenchRun(): Task[] {
+    teardownCore();
+    console.log(`[CodSpeed] Done running ${this.bench.tasks.length} benches.`);
+    return this.bench.tasks;
+  }
+
+  protected wrapWithInstrumentHooks<T>(fn: () => T, uri: string): T {
+    InstrumentHooks.startBenchmark();
+    const result = fn();
+    InstrumentHooks.stopBenchmark();
+    InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    return result;
+  }
+
+  protected async wrapWithInstrumentHooksAsync(
+    fn: Fn,
+    uri: string
+  ): Promise<unknown> {
+    InstrumentHooks.startBenchmark();
+    const result = await fn();
+    InstrumentHooks.stopBenchmark();
+    InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    return result;
+  }
+
+  protected abstract getModeName(): string;
+  protected abstract runTaskAsync(task: Task, uri: string): Promise<void>;
+  protected abstract runTaskSync(task: Task, uri: string): void;
+  protected abstract finalizeAsyncRun(): Task[];
+  protected abstract finalizeSyncRun(): Task[];
+
+  public setupBenchMethods(): void {
+    this.bench.run = async () => {
+      this.setupBenchRun();
+
+      for (const task of this.bench.tasks) {
+        const uri = this.getTaskUri(task);
+        await this.runTaskAsync(task, uri);
+      }
+
+      return this.finalizeAsyncRun();
+    };
+
+    this.bench.runSync = () => {
+      this.setupBenchRun();
+
+      for (const task of this.bench.tasks) {
+        const uri = this.getTaskUri(task);
+        this.runTaskSync(task, uri);
+      }
+
+      return this.finalizeSyncRun();
+    };
+  }
+}