feat: to fixup tinybench marker clean up

Author: GuillaumeLagrange

packages/tinybench-plugin/src/analysis.ts

@@ -3,6 +3,7 @@ import {
   InstrumentHooks,
   mongoMeasurement,
   optimizeFunction,
+  optimizeFunctionSync,
   wrapWithRootFrame,
   wrapWithRootFrameSync,
 } from "@codspeed/core";
@@ -40,7 +41,14 @@ class AnalysisBenchRunner extends BaseBenchRunner {
     await mongoMeasurement.start(uri);
 
     global.gc?.();
-    await this.wrapWithInstrumentHooksAsync(wrapWithRootFrame(fn), uri);
+    const rootFrame = wrapWithRootFrame(fn);
+    InstrumentHooks.startBenchmark();
+    try {
+      await rootFrame();
+    } finally {
+      InstrumentHooks.stopBenchmark();
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    }
 
     await mongoMeasurement.stop(uri);
     await fnOpts?.afterEach?.call(task, "run");
@@ -53,9 +61,24 @@ class AnalysisBenchRunner extends BaseBenchRunner {
     const { fnOpts, fn } = this.getTaskData(task);
 
     fnOpts?.beforeAll?.call(task, "run");
+
+    optimizeFunctionSync(async () => {
+      fnOpts?.beforeEach?.call(task, "run");
+      await fn();
+      fnOpts?.afterEach?.call(task, "run");
+    });
+
     fnOpts?.beforeEach?.call(task, "run");
 
-    this.wrapWithInstrumentHooks(wrapWithRootFrameSync(fn), uri);
+    global.gc?.();
+    const rootFrame = wrapWithRootFrameSync(fn);
+    InstrumentHooks.startBenchmark();
+    try {
+      rootFrame();
+    } finally {
+      InstrumentHooks.stopBenchmark();
+      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
+    }
 
     fnOpts?.afterEach?.call(task, "run");
     fnOpts?.afterAll?.call(task, "run");

packages/tinybench-plugin/src/shared.ts

@@ -1,12 +1,5 @@
-import {
-  getInstrumentMode,
-  InstrumentHooks,
-  MARKER_TYPE_BENCHMARK_END,
-  MARKER_TYPE_BENCHMARK_START,
-  setupCore,
-  teardownCore,
-} from "@codspeed/core";
-import { Bench, Fn, Task } from "tinybench";
+import { setupCore, teardownCore } from "@codspeed/core";
+import { Bench, Task } from "tinybench";
 import { CapturedTaskData, getTaskData } from "./taskData";
 import { getTaskUri } from "./uri";
 
@@ -56,57 +49,12 @@ export abstract class BaseBenchRunner {
     return this.bench.tasks;
   }
 
-  protected wrapWithInstrumentHooks<T>(fn: () => T, uri: string): T {
-    const runStart = this.openInstrumentWindow();
-    try {
-      return fn();
-    } finally {
-      this.closeInstrumentWindow(uri, runStart);
-    }
-  }
-
-  protected async wrapWithInstrumentHooksAsync(
-    fn: Fn,
-    uri: string,
-  ): Promise<unknown> {
-    const runStart = this.openInstrumentWindow();
-    try {
-      return await fn();
-    } finally {
-      this.closeInstrumentWindow(uri, runStart);
-    }
-  }
-
-  protected openInstrumentWindow(): bigint {
-    InstrumentHooks.startBenchmark();
-    return InstrumentHooks.currentTimestamp();
-  }
-
-  protected closeInstrumentWindow(uri: string, runStart: bigint): void {
-    const runEnd = InstrumentHooks.currentTimestamp();
-    InstrumentHooks.stopBenchmark();
-    InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-    this.sendBenchmarkMarkers(runStart, runEnd);
-  }
-
   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[];
 
-  private sendBenchmarkMarkers(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();

packages/tinybench-plugin/src/walltime.ts

@@ -1,5 +1,8 @@
 import {
   calculateQuantiles,
+  InstrumentHooks,
+  MARKER_TYPE_BENCHMARK_END,
+  MARKER_TYPE_BENCHMARK_START,
   mongoMeasurement,
   msToNs,
   msToS,
@@ -54,20 +57,39 @@ class WalltimeBenchRunner extends BaseBenchRunner {
     opts.setup = (task, mode) => {
       const setupResult = userSetup(task, mode);
       if (mode === "run") {
-        this.runStart = this.openInstrumentWindow();
+        InstrumentHooks.startBenchmark();
+        this.runStart = InstrumentHooks.currentTimestamp();
       }
       return setupResult;
     };
 
     opts.teardown = (task, mode) => {
       if (mode === "run" && task) {
-        this.closeInstrumentWindow(this.getTaskUri(task), this.runStart!);
+        const runEnd = InstrumentHooks.currentTimestamp();
+        InstrumentHooks.stopBenchmark();
+        InstrumentHooks.setExecutedBenchmark(
+          process.pid,
+          this.getTaskUri(task),
+        );
+        if (this.runStart !== null) {
+          this.sendBenchmarkMarkers(this.runStart, runEnd);
+        }
+
         this.runStart = null;
       }
       return userTeardown(task, mode);
     };
   }
 
+  private sendBenchmarkMarkers(runStart: bigint, runEnd: bigint): void {
+    InstrumentHooks.addMarker(
+      process.pid,
+      MARKER_TYPE_BENCHMARK_START,
+      runStart,
+    );
+    InstrumentHooks.addMarker(process.pid, MARKER_TYPE_BENCHMARK_END, runEnd);
+  }
+
   protected async runTaskAsync(task: Task, uri: string): Promise<void> {
     // run the warmup of the task right before its actual run
     if (getBenchOptions(this.bench).warmup) {