fixup! feat(tinybench-plugin): bump tinybench and add walltime support

Author: GuillaumeLagrange

packages/core/src/index.ts

@@ -10,24 +10,26 @@ export const isBound = native_core.isBound;
 
 export const mongoMeasurement = new MongoMeasurement();
 
-export enum MeasurementMode {
-  Instrumentation = "instrumentation",
-  WallTime = "walltime",
-}
+type CodSpeedRunnerMode = "disabled" | "instrumented" | "walltime";
 
-export function getMeasurementMode(): MeasurementMode {
+export function getCodspeedRunnerMode(): CodSpeedRunnerMode {
   const isCodSpeedEnabled = process.env.CODSPEED_ENV !== undefined;
-  if (isCodSpeedEnabled) {
-    // If CODSPEED_ENV is set, check CODSPEED_RUNNER_MODE
-    if (process.env.CODSPEED_RUNNER_MODE === "walltime") {
-      return MeasurementMode.WallTime;
-    } else {
-      return MeasurementMode.Instrumentation;
-    }
+  if (!isCodSpeedEnabled) {
+    return "disabled";
+  }
+
+  // If CODSPEED_ENV is set, check CODSPEED_RUNNER_MODE
+  const codspeedRunnerMode = process.env.CODSPEED_RUNNER_MODE;
+  if (codspeedRunnerMode === "instrumentation") {
+    return "instrumented";
+  } else if (codspeedRunnerMode === "walltime") {
+    return "walltime";
   }
 
-  // Default to walltime mode when CODSPEED_ENV is not set
-  return MeasurementMode.WallTime;
+  console.warn(
+    `Unknown codspeed runner mode: ${codspeedRunnerMode}, defaulting to disabled`
+  );
+  return "disabled";
 }
 
 export const setupCore = () => {

packages/tinybench-plugin/src/index.ts

@@ -1,4 +1,5 @@
 import {
+  getCodspeedRunnerMode,
   getGitDir,
   Measurement,
   mongoMeasurement,
@@ -11,53 +12,33 @@ import { get as getStackTrace } from "stack-trace";
 import { Bench } from "tinybench";
 import { fileURLToPath } from "url";
 import { runInstrumentedBench } from "./instrumented";
+import { getOrCreateUriMap } from "./uri";
 import { runWalltimeBench } from "./walltime";
 
 tryIntrospect();
 
-// Store URI mapping externally since fnOpts is private
-export const taskUriMap = new WeakMap<Bench, Map<string, string>>();
-
 export function withCodSpeed(bench: Bench): Bench {
-  // Check if CODSPEED_ENV is defined
-  if (!process.env.CODSPEED_ENV) {
-    // If CODSPEED_ENV is not defined, return bench unchanged (run as normal tinybench)
+  const codspeedRunnerMode = getCodspeedRunnerMode();
+  if (codspeedRunnerMode === "disabled") {
     return bench;
   }
 
   const rootCallingFile = getCallingFile();
 
-  // Initialize URI mapping for this bench instance
-  if (!taskUriMap.has(bench)) {
-    taskUriMap.set(bench, new Map());
-  }
-  const uriMap = taskUriMap.get(bench)!;
-
-  // Setup URI generation for tasks
+  // Compute and register URI for bench
+  const uriMap = getOrCreateUriMap(bench);
   const rawAdd = bench.add;
   bench.add = (name, fn, opts?) => {
     const callingFile = getCallingFile();
     const uri = `${callingFile}::${name}`;
-    // Store URI mapping
     uriMap.set(name, uri);
     return rawAdd.bind(bench)(name, fn, opts);
   };
 
-  // Apply the appropriate measurement strategy based on CODSPEED_RUNNER_MODE
-  const runnerMode = process.env.CODSPEED_RUNNER_MODE;
-
-  if (runnerMode === "instrumentation") {
+  if (codspeedRunnerMode === "instrumented") {
     runInstrumentedBench(bench, rootCallingFile);
-  } else if (runnerMode === "walltime") {
+  } else if (codspeedRunnerMode === "walltime") {
     runWalltimeBench(bench, rootCallingFile);
-  } else {
-    const rawRun = bench.run;
-    bench.run = async () => {
-      console.warn(
-        `[CodSpeed] CODSPEED_ENV is defined but CODSPEED_RUNNER_MODE="${runnerMode}" is not recognized, falling back to tinybench`
-      );
-      return await rawRun.bind(bench)();
-    };
   }
 
   return bench;

packages/tinybench-plugin/src/instrumented.ts

@@ -5,20 +5,11 @@ import {
   setupCore,
   teardownCore,
 } from "@codspeed/core";
-import { Bench } from "tinybench";
-import { taskUriMap } from "./index";
+import { Bench, Fn, FnOptions } from "tinybench";
+import { getTaskUri } from "./uri";
 
 declare const __VERSION__: string;
 
-function getTaskUri(
-  bench: Bench,
-  taskName: string,
-  rootCallingFile: string
-): string {
-  const uriMap = taskUriMap.get(bench);
-  return uriMap?.get(taskName) || `${rootCallingFile}::${taskName}`;
-}
-
 export function runInstrumentedBench(
   bench: Bench,
   rootCallingFile: string
@@ -32,16 +23,16 @@ export function runInstrumentedBench(
     for (const task of bench.tasks) {
       const uri = getTaskUri(bench, task.name, rootCallingFile);
 
-      // Access private fnOpts to get hooks
-      const fnOpts = (task as any).fnOpts;
+      // Access private fields
+      const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
       // Call beforeAll hook if it exists
       await fnOpts?.beforeAll?.call(task, "run");
 
       // run optimizations
       await optimizeFunction(async () => {
         await fnOpts?.beforeEach?.call(task, "run");
-        await (task as any).fn(); // Access private fn property
+        await fn();
         await fnOpts?.afterEach?.call(task, "run");
       });
 
@@ -52,7 +43,7 @@ export function runInstrumentedBench(
       global.gc?.();
       await (async function __codspeed_root_frame__() {
         Measurement.startInstrumentation();
-        await (task as any).fn(); // Access private fn property
+        await fn();
         Measurement.stopInstrumentation(uri);
       })();
       await mongoMeasurement.stop(uri);

packages/tinybench-plugin/src/uri.ts

@@ -0,0 +1,20 @@
+import { Bench } from "tinybench";
+
+// Store URI mapping externally since fnOpts is private
+export const taskUriMap = new WeakMap<Bench, Map<string, string>>();
+
+export function getTaskUri(
+  bench: Bench,
+  taskName: string,
+  rootCallingFile: string
+): string {
+  const uriMap = taskUriMap.get(bench);
+  return uriMap?.get(taskName) || `${rootCallingFile}::${taskName}`;
+}
+
+export function getOrCreateUriMap(bench: Bench) {
+  if (!taskUriMap.has(bench)) {
+    taskUriMap.set(bench, new Map());
+  }
+  return taskUriMap.get(bench)!;
+}

packages/tinybench-plugin/src/walltime.ts

@@ -6,53 +6,10 @@ import {
   type BenchmarkStats,
 } from "@codspeed/core";
 import { Bench, TaskResult } from "tinybench";
-import { taskUriMap } from "./index";
+import { getTaskUri } from "./uri";
 
 declare const __VERSION__: string;
 
-function getTaskUri(
-  bench: Bench,
-  taskName: string,
-  rootCallingFile: string
-): string {
-  const uriMap = taskUriMap.get(bench);
-  return uriMap?.get(taskName) || `${rootCallingFile}::${taskName}`;
-}
-
-function convertTinybenchResultToBenchmarkStats(
-  result: TaskResult,
-  warmupIterations: number
-): BenchmarkStats {
-  // All tinybench times are in milliseconds, convert to nanoseconds
-  const ms_to_ns = (ms: number) => ms * 1_000_000;
-
-  const { min, max, mean, sd, samples } = result.latency;
-
-  // Get individual sample times in nanoseconds and sort them
-  const sortedTimesNs = samples.map(ms_to_ns).sort((a, b) => a - b);
-  const meanNs = ms_to_ns(mean);
-  const stdevNs = ms_to_ns(sd);
-
-  const { q1_ns, q3_ns, median_ns, iqr_outlier_rounds, stdev_outlier_rounds } =
-    calculateQuantiles({ meanNs, stdevNs, sortedTimesNs });
-
-  return {
-    min_ns: ms_to_ns(min),
-    max_ns: ms_to_ns(max),
-    mean_ns: meanNs,
-    stdev_ns: stdevNs,
-    q1_ns,
-    median_ns,
-    q3_ns,
-    total_time: result.totalTime / 1_000, // convert from ms to seconds
-    iter_per_round: sortedTimesNs.length,
-    rounds: 1, // Tinybench only runs one round
-    iqr_outlier_rounds,
-    stdev_outlier_rounds,
-    warmup_iters: warmupIterations,
-  };
-}
-
 export function runWalltimeBench(bench: Bench, rootCallingFile: string): void {
   bench.run = async () => {
     console.log(
@@ -125,3 +82,37 @@ export function runWalltimeBench(bench: Bench, rootCallingFile: string): void {
     return results;
   };
 }
+
+function convertTinybenchResultToBenchmarkStats(
+  result: TaskResult,
+  warmupIterations: number
+): BenchmarkStats {
+  // All tinybench times are in milliseconds, convert to nanoseconds
+  const ms_to_ns = (ms: number) => ms * 1_000_000;
+
+  const { min, max, mean, sd, samples } = result.latency;
+
+  // Get individual sample times in nanoseconds and sort them
+  const sortedTimesNs = samples.map(ms_to_ns).sort((a, b) => a - b);
+  const meanNs = ms_to_ns(mean);
+  const stdevNs = ms_to_ns(sd);
+
+  const { q1_ns, q3_ns, median_ns, iqr_outlier_rounds, stdev_outlier_rounds } =
+    calculateQuantiles({ meanNs, stdevNs, sortedTimesNs });
+
+  return {
+    min_ns: ms_to_ns(min),
+    max_ns: ms_to_ns(max),
+    mean_ns: meanNs,
+    stdev_ns: stdevNs,
+    q1_ns,
+    median_ns,
+    q3_ns,
+    total_time: result.totalTime / 1_000, // convert from ms to seconds
+    iter_per_round: sortedTimesNs.length,
+    rounds: 1, // Tinybench only runs one round
+    iqr_outlier_rounds,
+    stdev_outlier_rounds,
+    warmup_iters: warmupIterations,
+  };
+}