chore(tinybench): emit benchmark markers

Author: not-matthias

packages/tinybench-plugin/src/shared.ts

@@ -1,4 +1,11 @@
-import { InstrumentHooks, setupCore, teardownCore } from "@codspeed/core";
+import {
+  getInstrumentMode,
+  InstrumentHooks,
+  MARKER_TYPE_BENCHMARK_END,
+  MARKER_TYPE_BENCHMARK_START,
+  setupCore,
+  teardownCore,
+} from "@codspeed/core";
 import { Bench, Fn, Task } from "tinybench";
 import { getTaskUri } from "./uri";
 
@@ -63,25 +70,43 @@ export abstract class BaseBenchRunner {
   protected abstract finalizeAsyncRun(): Task[];
   protected abstract finalizeSyncRun(): Task[];
 
+  private sendRunMarkers(runStart: bigint, runEnd: bigint): void {
+    if (getInstrumentMode() !== "walltime") {
+      return;
+    }
+    InstrumentHooks.addMarker(
+      process.pid,
+      MARKER_TYPE_BENCHMARK_START,
+      runStart,
+    );
+    InstrumentHooks.addMarker(process.pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+  }
+
   public setupBenchMethods(): void {
     this.bench.run = async () => {
       this.setupBenchRun();
 
+      const runStart = InstrumentHooks.currentTimestamp();
       for (const task of this.bench.tasks) {
         const uri = this.getTaskUri(task);
         await this.runTaskAsync(task, uri);
       }
+      const runEnd = InstrumentHooks.currentTimestamp();
+      this.sendRunMarkers(runStart, runEnd);
 
       return this.finalizeAsyncRun();
     };
 
     this.bench.runSync = () => {
       this.setupBenchRun();
 
+      const runStart = InstrumentHooks.currentTimestamp();
       for (const task of this.bench.tasks) {
         const uri = this.getTaskUri(task);
         this.runTaskSync(task, uri);
       }
+      const runEnd = InstrumentHooks.currentTimestamp();
+      this.sendRunMarkers(runStart, runEnd);
 
       return this.finalizeSyncRun();
     };

packages/tinybench-plugin/tests/index.integ.test.ts

@@ -16,6 +16,8 @@ const mockCore = vi.hoisted(() => {
       startBenchmark: vi.fn(),
       stopBenchmark: vi.fn(),
       setExecutedBenchmark: vi.fn(),
+      currentTimestamp: vi.fn().mockReturnValue(0n),
+      addMarker: vi.fn(),
     },
     optimizeFunction: vi
       .fn()

packages/vitest-plugin/src/walltime/index.ts

@@ -5,7 +5,6 @@ import {
   setupCore,
   writeWalltimeResults,
 } from "@codspeed/core";
-import { Fn } from "tinybench";
 import {
   RunnerTaskEventPack,
   RunnerTaskResultPack,
@@ -78,37 +77,18 @@ export class WalltimeRunner extends NodeBenchmarkRunner {
     };
 
     tinybench.Task.prototype.run = async function () {
-      const { fn } = this as { fn: Fn };
       const suiteUri = getSuiteUri();
 
-      const finishRound = (roundStart: bigint): void => {
-        const roundEnd = InstrumentHooks.currentTimestamp();
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, roundStart);
-        InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, roundEnd);
-      };
-
-      function __codspeed_root_frame__(): unknown {
-        const roundStart = InstrumentHooks.currentTimestamp();
-        const result = fn();
-        if (
-          result !== null &&
-          typeof result === "object" &&
-          typeof (result as PromiseLike<unknown>).then === "function"
-        ) {
-          return (result as PromiseLike<unknown>).then((value) => {
-            finishRound(roundStart);
-            return value;
-          });
-        }
-        finishRound(roundStart);
-        return result;
-      }
-      (this as { fn: Fn }).fn = __codspeed_root_frame__ as Fn;
-
       InstrumentHooks.startBenchmark();
+      const runStart = InstrumentHooks.currentTimestamp();
       await originalRun.call(this);
+      const runEnd = InstrumentHooks.currentTimestamp();
       InstrumentHooks.stopBenchmark();
 
+      // Emit a single marker pair covering the whole measurement run
+      InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_START, runStart);
+      InstrumentHooks.addMarker(pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+
       // Look up the URI by task name
       const uri = `${suiteUri}::${this.name}`;
       InstrumentHooks.setExecutedBenchmark(pid, uri);