refactor(tinybench-plugin): commonize logic between run and runSync

Author: GuillaumeLagrange

packages/tinybench-plugin/src/instrumented.ts

@@ -4,7 +4,7 @@ import {
   optimizeFunction,
   teardownCore,
 } from "@codspeed/core";
-import { Bench, Fn, FnOptions } from "tinybench";
+import { Bench, Fn, FnOptions, Task } from "tinybench";
 import { getTaskUri } from "./uri";
 
 declare const __VERSION__: string;
@@ -13,103 +13,89 @@ export function runInstrumentedBench(
   bench: Bench,
   rootCallingFile: string
 ): void {
-  bench.run = async () => {
-    console.log(
-      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
-    );
-
-    for (const task of bench.tasks) {
-      const uri = getTaskUri(bench, task.name, rootCallingFile);
+  const runTaskAsync = async (task: Task, uri: string): Promise<void> => {
+    const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
-      // Access private fields
-      const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
+    await fnOpts?.beforeAll?.call(task, "run");
+    await optimizeFunction(async () => {
+      await fnOpts?.beforeEach?.call(task, "run");
+      await fn();
+      await fnOpts?.afterEach?.call(task, "run");
+    });
+    await fnOpts?.beforeEach?.call(task, "run");
+    await mongoMeasurement.start(uri);
 
-      // Call beforeAll hook if it exists
-      await fnOpts?.beforeAll?.call(task, "run");
+    await (async function __codspeed_root_frame__() {
+      global.gc?.();
+      InstrumentHooks.startBenchmark();
+      await fn();
+      InstrumentHooks.stopBenchmark();
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    })();
+
+    await mongoMeasurement.stop(uri);
+    await fnOpts?.afterEach?.call(task, "run");
+    await fnOpts?.afterAll?.call(task, "run");
+  };
 
-      // run optimizations
-      await optimizeFunction(async () => {
-        await fnOpts?.beforeEach?.call(task, "run");
-        await fn();
-        await fnOpts?.afterEach?.call(task, "run");
-      });
+  // Sync task runner
+  const runTaskSync = (task: Task, uri: string): void => {
+    const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
 
-      // run instrumented benchmark
-      await fnOpts?.beforeEach?.call(task, "run");
+    fnOpts?.beforeAll?.call(task, "run");
+    fnOpts?.beforeEach?.call(task, "run");
 
-      await mongoMeasurement.start(uri);
+    (function __codspeed_root_frame__() {
       global.gc?.();
-      await (async function __codspeed_root_frame__() {
-        InstrumentHooks.startBenchmark();
-        await fn();
-        InstrumentHooks.stopBenchmark();
-        InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-      })();
-      await mongoMeasurement.stop(uri);
-
-      await fnOpts?.afterEach?.call(task, "run");
+      InstrumentHooks.startBenchmark();
+      fn();
+      InstrumentHooks.stopBenchmark();
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    })();
+
+    fnOpts?.afterEach?.call(task, "run");
+    fnOpts?.afterAll?.call(task, "run");
+  };
 
-      await fnOpts?.afterAll?.call(task, "run");
+  bench.run = async () => {
+    logStart();
 
-      // print results
-      console.log(
-        `    ✔ ${
-          InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
-        } ${uri}`
-      );
+    for (const task of bench.tasks) {
+      const uri = getTaskUri(bench, task.name, rootCallingFile);
+      await runTaskAsync(task, uri);
+      logTaskCompletion(uri);
     }
 
-    teardownCore();
-    console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);
-    return bench.tasks;
+    return logEnd();
   };
 
   bench.runSync = () => {
-    console.log(
-      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
-    );
+    logStart();
 
     for (const task of bench.tasks) {
       const uri = getTaskUri(bench, task.name, rootCallingFile);
+      runTaskSync(task, uri);
+      logTaskCompletion(uri);
+    }
 
-      // Access private fields
-      const { fnOpts, fn } = task as unknown as { fnOpts?: FnOptions; fn: Fn };
-
-      // Call beforeAll hook if it exists
-      fnOpts?.beforeAll?.call(task, "run");
-
-      // run optimizations
-      optimizeFunction(async () => {
-        fnOpts?.beforeEach?.call(task, "run");
-        fn();
-        fnOpts?.afterEach?.call(task, "run");
-      });
+    return logEnd();
+  };
 
-      // run instrumented benchmark
-      fnOpts?.beforeEach?.call(task, "run");
+  const logStart = () => {
+    console.log(
+      `[CodSpeed] running with @codspeed/tinybench v${__VERSION__} (instrumented mode)`
+    );
+  };
 
-      // await mongoMeasurement.start(uri);
-      global.gc?.();
-      (function __codspeed_root_frame__() {
-        InstrumentHooks.startBenchmark();
-        fn();
-        InstrumentHooks.stopBenchmark();
-        InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-      })();
-      mongoMeasurement.stop(uri);
-
-      fnOpts?.afterEach?.call(task, "run");
-
-      fnOpts?.afterAll?.call(task, "run");
-
-      // print results
-      console.log(
-        `    ✔ ${
-          InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
-        } ${uri}`
-      );
-    }
+  const logTaskCompletion = (uri: string) => {
+    console.log(
+      `    ✔ ${
+        InstrumentHooks.isInstrumented() ? "Measured" : "Checked"
+      } ${uri}`
+    );
+  };
 
+  const logEnd = () => {
     teardownCore();
     console.log(`[CodSpeed] Done running ${bench.tasks.length} benches.`);
     return bench.tasks;