fix: cleanup handler for task after completion (#21)

* Assign and cleanup the task fn

* Bump rc version

Author: nazarhussain

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@chainsafe/benchmark",
-  "version": "1.1.0-rc.5",
+  "version": "1.1.0-rc.6",
   "repository": "[email protected]:chainsafe/benchmark.git",
   "author": "ChainSafe Systems",
   "license": "MIT",

src/benchmark/benchmarkFn.ts

@@ -1,6 +1,6 @@
 import fs from "node:fs";
 import path from "node:path";
-import {getCurrentSuite} from "@vitest/runner";
+import {getCurrentSuite, setFn} from "@vitest/runner";
 import {createChainable} from "@vitest/runner/utils";
 import {store} from "./globalState.js";
 import {BenchApi, BenchmarkOpts, BenchmarkRunOptsWithFn, PartialBy} from "../types.js";
@@ -29,6 +29,31 @@ export const bench: BenchApi = createBenchmarkFunction(function <T, T2>(
     timeout = minMs * 1.5;
   }
 
+  async function handler(): Promise<void> {
+    // Ensure bench id is unique
+    if (store.getResult(opts.id) && !opts.skip) {
+      throw Error(`test titles must be unique, duplicated: '${opts.id}'`);
+    }
+
+    // Persist full results if requested. dir is created in `beforeAll`
+    const benchmarkResultsCsvDir = process.env.BENCHMARK_RESULTS_CSV_DIR;
+    const persistRunsNs = Boolean(benchmarkResultsCsvDir);
+
+    const {result, runsNs} = await runBenchFn({...options, fn: benchTask}, persistRunsNs);
+
+    // Store result for:
+    // - to persist benchmark data latter
+    // - to render with the custom reporter
+    store.setResult(opts.id, result);
+
+    if (benchmarkResultsCsvDir) {
+      fs.mkdirSync(benchmarkResultsCsvDir, {recursive: true});
+      const filename = `${result.id}.csv`;
+      const filepath = path.join(benchmarkResultsCsvDir, filename);
+      fs.writeFileSync(filepath, runsNs.join("\n"));
+    }
+  }
+
   const task = currentSuite.task(opts.id, {
     skip: opts.skip ?? this.skip,
     only: opts.only ?? this.only,
@@ -38,32 +63,9 @@ export const bench: BenchApi = createBenchmarkFunction(function <T, T2>(
     meta: {
       "chainsafe/benchmark": true,
     },
-    async handler() {
-      // Ensure bench id is unique
-      if (store.getResult(opts.id) && !opts.skip) {
-        throw Error(`test titles must be unique, duplicated: '${opts.id}'`);
-      }
-
-      // Persist full results if requested. dir is created in `beforeAll`
-      const benchmarkResultsCsvDir = process.env.BENCHMARK_RESULTS_CSV_DIR;
-      const persistRunsNs = Boolean(benchmarkResultsCsvDir);
-
-      const {result, runsNs} = await runBenchFn({...options, fn: benchTask}, persistRunsNs);
-
-      // Store result for:
-      // - to persist benchmark data latter
-      // - to render with the custom reporter
-      store.setResult(opts.id, result);
-
-      if (benchmarkResultsCsvDir) {
-        fs.mkdirSync(benchmarkResultsCsvDir, {recursive: true});
-        const filename = `${result.id}.csv`;
-        const filepath = path.join(benchmarkResultsCsvDir, filename);
-        fs.writeFileSync(filepath, runsNs.join("\n"));
-      }
-    },
   });
 
+  setFn(task, handler);
   store.setOptions(task, opts);
 });
 
@@ -85,7 +87,7 @@ function coerceToOptsObj<T, T2>(
 
   if (typeof idOrOpts === "string") {
     if (!fn) throw Error("fn arg must be set");
-    opts = {id: idOrOpts, fn, threshold: optionsDefault.threshold};
+    opts = {id: idOrOpts, fn};
   } else {
     if (fn) {
       opts = {...idOrOpts, fn};

src/benchmark/runner.ts

@@ -6,6 +6,7 @@ import {
   VitestRunner,
   VitestRunnerConfig,
   VitestRunnerImportSource,
+  setFn,
 } from "@vitest/runner";
 import path from "node:path";
 import {Benchmark, BenchmarkOpts, BenchmarkResults} from "../types.js";
@@ -49,25 +50,37 @@ export class BenchmarkRunner implements VitestRunner {
     this.reporter.onTestStarted(task);
   }
 
-  onAfterRunTask(task: Task): void {
+  async onAfterRunTask(task: Task): Promise<void> {
     this.reporter.onTestFinished(task);
     store.removeOptions(task);
 
+    if (task.type === "test" || task.type === "custom") {
+      // Clear up the assigned handler to clean the memory
+      // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+      // @ts-expect-error
+      setFn(task, null);
+    }
+
     // To help maintain consistent memory usage patterns
     // we trigger garbage collection manually
     if (this.triggerGC && global.gc) {
       global.gc();
+      // Make sure the syn operation is off the event loop
+      await new Promise((resolve) => setTimeout(resolve, 0));
     }
   }
 
   onAfterRunFiles(files: File[]): void {
     this.reporter.onComplete(files);
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-unused-vars
-  async importFile(filepath: string, _source: VitestRunnerImportSource): Promise<void> {
+  async importFile(filepath: string, source: VitestRunnerImportSource): Promise<void> {
+    let url = filepath;
+    if (source === "setup") {
+      url = `${url}?key=${Date.now()}`;
+    }
     // TODO: Implement file caching mechanism later
-    await import(filepath);
+    await import(url);
   }
 
   async process(files: string[]): Promise<BenchmarkResults> {