fix(tinybench-plugin): fix profiler capturing tinybench's internal functions

Author: GuillaumeLagrange

CLAUDE.md

@@ -88,4 +88,10 @@ Based on the codebase analysis, to add stats access features:
 ## Repository Management Memories
 
 - Use pnpm instead of npm
-- To run tests in a package use moon <package-name>:test
+- The monorepo uses [Turborepo](https://turborepo.com). Run a task for a single
+  package with a filter, using the package's full name (e.g. `@codspeed/tinybench-plugin`):
+  - Test: `pnpm turbo run test --filter=<package-name>`
+  - Build: `pnpm turbo run build --filter=<package-name>`
+  - Typecheck: `pnpm turbo run typecheck --filter=<package-name>`
+  - Lint: `pnpm turbo run lint --filter=<package-name>`
+  - Run a task across all packages by omitting `--filter` (e.g. `pnpm turbo run build`).

packages/tinybench-plugin/src/shared.ts

@@ -46,34 +46,38 @@ export abstract class BaseBenchRunner {
   }
 
   protected wrapWithInstrumentHooks<T>(fn: () => T, uri: string): T {
-    InstrumentHooks.startBenchmark();
-    const runStart = InstrumentHooks.currentTimestamp();
+    const runStart = this.openInstrumentWindow();
     try {
       return fn();
     } finally {
-      const runEnd = InstrumentHooks.currentTimestamp();
-      InstrumentHooks.stopBenchmark();
-      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-      this.sendBenchmarkMarkers(runStart, runEnd);
+      this.closeInstrumentWindow(uri, runStart);
     }
   }
 
   protected async wrapWithInstrumentHooksAsync(
     fn: Fn,
     uri: string,
   ): Promise<unknown> {
-    InstrumentHooks.startBenchmark();
-    const runStart = InstrumentHooks.currentTimestamp();
+    const runStart = this.openInstrumentWindow();
     try {
       return await fn();
     } finally {
-      const runEnd = InstrumentHooks.currentTimestamp();
-      InstrumentHooks.stopBenchmark();
-      InstrumentHooks.setExecutedBenchmark(process.pid, uri);
-      this.sendBenchmarkMarkers(runStart, runEnd);
+      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;

packages/tinybench-plugin/src/walltime.ts

@@ -9,24 +9,60 @@ import {
   type BenchmarkStats,
   type Benchmark as CodspeedBenchmark,
 } from "@codspeed/core";
-import { Bench, Fn, Task, TaskResult } from "tinybench";
+import { Bench, Fn, Hook, Task, TaskResult } from "tinybench";
 import { BaseBenchRunner } from "./shared";
 
 export function setupCodspeedWalltimeBench(
   bench: Bench,
   rootCallingFile: string,
 ): void {
   const runner = new WalltimeBenchRunner(bench, rootCallingFile);
+  runner.installInstrumentHooks();
   runner.setupBenchMethods();
 }
 
 class WalltimeBenchRunner extends BaseBenchRunner {
   private codspeedBenchmarks: CodspeedBenchmark[] = [];
 
+  // Carries the window start timestamp from the setup hook to the teardown
+  // hook. Tasks run strictly sequentially, so a single field is enough.
+  private runStart: bigint | null = null;
+
   protected getModeName(): string {
     return "walltime mode";
   }
 
+  /**
+   * Drive the instrumentation window from the task's setup/teardown hooks so it
+   * brackets only tinybench's measured loop, excluding warmup and the
+   * statistics computation (`processRunResult`) that surround it. The user's
+   * own hooks are preserved and still run in their original order relative to
+   * the work under test.
+   */
+  public installInstrumentHooks(): void {
+    // `bench.opts` is typed `Readonly`, but tinybench mutates it at runtime and
+    // always resolves `setup`/`teardown` to (at least) a no-op default.
+    const opts = this.bench.opts as { setup: Hook; teardown: Hook };
+    const userSetup = opts.setup;
+    const userTeardown = opts.teardown;
+
+    opts.setup = (task, mode) => {
+      const setupResult = userSetup(task, mode);
+      if (mode === "run") {
+        this.runStart = this.openInstrumentWindow();
+      }
+      return setupResult;
+    };
+
+    opts.teardown = (task, mode) => {
+      if (mode === "run" && task) {
+        this.closeInstrumentWindow(this.getTaskUri(task), this.runStart!);
+        this.runStart = null;
+      }
+      return userTeardown(task, mode);
+    };
+  }
+
   protected async runTaskAsync(task: Task, uri: string): Promise<void> {
     // Override the function under test to add a static frame
     this.wrapTaskFunction(task, true);
@@ -37,21 +73,22 @@ class WalltimeBenchRunner extends BaseBenchRunner {
     }
 
     await mongoMeasurement.start(uri);
-    await this.wrapWithInstrumentHooksAsync(() => task.run(), uri);
+    await task.run();
     await mongoMeasurement.stop(uri);
 
     this.registerCodspeedBenchmarkFromTask(task);
   }
 
-  protected runTaskSync(task: Task, uri: string): void {
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  protected runTaskSync(task: Task, _uri: string): void {
     // Override the function under test to add a static frame
     this.wrapTaskFunction(task, false);
 
     if (this.bench.opts.warmup) {
       task.warmup();
     }
 
-    this.wrapWithInstrumentHooks(() => task.runSync(), uri);
+    task.runSync();
 
     this.registerCodspeedBenchmarkFromTask(task);
   }