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,
@@ -23,7 +22,6 @@ export function withCodSpeed(bench: Bench): Bench {
   if (codspeedRunnerMode === "disabled") {
     return bench;
   }
-  setupCore();
 
   const rootCallingFile = getCallingFile();
 

packages/tinybench-plugin/src/instrumented.ts

@@ -2,18 +2,34 @@ import {
   InstrumentHooks,
   mongoMeasurement,
   optimizeFunction,
-  teardownCore,
 } from "@codspeed/core";
 import { Bench, Fn, FnOptions, Task } from "tinybench";
-import { getTaskUri } from "./uri";
-
-declare const __VERSION__: string;
+import { BaseBenchRunner } from "./shared";
+
+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__() {
+        global.gc?.();
+        await fn();
+      };
+    } else {
+      return function __codspeed_root_frame__() {
+        global.gc?.();
+        fn();
+      };
+    }
+  }
 
-export function runInstrumentedBench(
-  bench: Bench,
-  rootCallingFile: string
-): void {
-  const runTaskAsync = async (task: Task, uri: string): Promise<void> => {
+  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 +41,48 @@ 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);
-    })();
+    await this.wrapWithHooksAsync(
+      this.wrapFunctionWithFrame(fn, true) as () => Promise<void>,
+      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.wrapWithHooks(
+      this.wrapFunctionWithFrame(fn, false) as () => void,
+      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);
-    }
-
-    return logEnd();
-  };
+    this.logTaskCompletion(uri, this.taskCompletionMessage());
+  }
 
-  const logStart = () => {
-    console.log(
-      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
-    );
-  };
+  protected finalizeAsyncRun(): Task[] {
+    return this.finalizeBenchRun();
+  }
 
-  const logTaskCompletion = (uri: string) => {
-    console.log(
-      `    ✔ ${
-        InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
-      } ${uri}`
-    );
-  };
+  protected finalizeSyncRun(): Task[] {
+    return this.finalizeBenchRun();
+  }
+}
 
-  const logEnd = () => {
-    teardownCore();
-    console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);
-    return bench.tasks;
-  };
+export function runInstrumentedBench(
+  bench: Bench,
+  rootCallingFile: string
+): void {
+  const runner = new InstrumentedBenchRunner(bench, rootCallingFile);
+  runner.setupBenchMethods();
 }

packages/tinybench-plugin/src/shared.ts

@@ -0,0 +1,89 @@
+import { InstrumentHooks, setupCore, teardownCore } from "@codspeed/core";
+import { Bench, 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 wrapWithHooks<T>(fn: () => T, uri: string): T {
+    InstrumentHooks.startBenchmark();
+    const result = fn();
+    InstrumentHooks.stopBenchmark();
+    InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    return result;
+  }
+
+  protected async wrapWithHooksAsync<T>(
+    fn: () => Promise<T>,
+    uri: string
+  ): Promise<T> {
+    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();
+    };
+  }
+}