feat: add support for embedded runner

README.md

@@ -9,10 +9,10 @@ This tooling provides both a easy to use runner for benchmarking and easy integr
 
 ## Quick start
 
-Create a test mocha test file but use `itBench` instead of `it`
+Create a benchmark test file but use `itBench` instead of `it`
 
 ```ts
-import {itBench, setBenchOpts} from "../../src";
+import {itBench, setBenchOpts, describe} from "../../src";
 
 describe("Sum array benchmark", () => {
   itBench("sum array with reduce", () => {
@@ -21,7 +21,7 @@ describe("Sum array benchmark", () => {
 });
 ```
 
-Then run the CLI, compatible with all mocha options.
+Then run the CLI.
 
 ```
 benchmark 'test/perf/**/*.perf.ts' --local
@@ -36,7 +36,7 @@ Inspect benchmark results in the terminal
 
 ## How does it work?
 
-This tool is a CLI wrapper around mocha, example usage:
+This tool is a CLI tool, example usage:
 
 ```
 benchmark 'test/perf/**/*.perf.ts' --s3
@@ -47,9 +47,7 @@ The above command will:
 - Read benchmark history from the specified provider (AWS S3)
 - Figure out the prev benchmark based on your option (defaults to latest commit in main branch)
 - Run benchmark comparing with previous
-  - Runs mocha programatically against the file globs
-  - Collect benchmark data in-memory while streaming results with a familiar mocha reporter
-  - Note: also runs any test that would regularly be run with mocha
+  - Collect benchmark data in-memory while streaming results
 - Add result to benchmark history and persist them to the specified provider (AWS S3)
 - If in CI, post a PR or commit comment with an expandable summary of the benchmark results comparision
 - If a performance regression was detected, exit 1

bin/index.cjs

@@ -1,3 +1,3 @@
 #!/usr/bin/env node
 
-require('../lib/cjs/cli.js');
+require("../lib/cjs/cli/cli.js");

package.json

@@ -33,14 +33,13 @@
     "test:unit": "mocha test/unit/**/*.test.ts",
     "lint": "eslint --color src/ test/",
     "prepublishOnly": "yarn build",
-    "benchmark": "node --loader ts-node/esm ./src/cli.ts 'test/perf/**/*.test.ts'",
+    "benchmark": "node --loader ts-node/esm ./src/cli/cli.ts 'test/perf/**/*.test.ts'",
     "writeDocs": "node --loader ts-node/esm scripts/writeOptionsMd.ts"
   },
   "devDependencies": {
     "@types/chai": "^4.2.19",
     "@types/mocha": "^10.0.9",
     "@types/node": "^18.15.3",
-    "@types/rimraf": "^3.0.0",
     "@types/yargs": "^17.0.33",
     "chai": "^4.5.0",
     "dotenv": "^10.0.0",
@@ -52,21 +51,24 @@
     "eslint-config-prettier": "^9.1.0",
     "mocha": "^10.8.2",
     "prettier": "^3.4.0",
-    "rimraf": "^3.0.2",
+    "rimraf": "^5.0.10",
     "ts-node": "^10.9.2",
     "typescript": "^5.6.3",
     "typescript-eslint": "^8.16.0"
   },
   "dependencies": {
     "@actions/cache": "^1.0.7",
     "@actions/github": "^5.0.0",
+    "@vitest/runner": "^2.1.6",    
     "ajv": "^8.6.0",
     "aws-sdk": "^2.932.0",
     "csv-parse": "^4.16.0",
     "csv-stringify": "^5.6.2",
-    "yargs": "^17.7.2"
+    "glob": "^10.4.5",
+    "yargs": "^17.7.2",
+    "log-symbols": "^7.0.0"
   },
-  "peerDependencies": {
-    "mocha": ">10.0.0"
+  "resolutions": {
+    "lru-cache": "10.4.3"
   }
 }

scripts/writeOptionsMd.ts

@@ -1,4 +1,4 @@
-import {options} from "../src/options.js";
+import {options} from "../src/cli/options.js";
 
 const sections: string[] = [];
 

src/benchmark/benchmarkFn.ts

@@ -0,0 +1,127 @@
+import fs from "node:fs";
+import path from "node:path";
+import {getCurrentSuite} from "@vitest/runner";
+import {createChainable} from "@vitest/runner/utils";
+import {store} from "./globalState.js";
+import {BenchApi, BenchmarkOpts, BenchmarkRunOptsWithFn, PartialBy} from "../types.js";
+import {runBenchFn} from "./runBenchmarkFn.js";
+import {optionsDefault} from "../cli/options.js";
+
+export const bench: BenchApi = createBenchmarkFunction(function <T, T2>(
+  this: Record<"skip" | "only", boolean | undefined>,
+  idOrOpts: string | PartialBy<BenchmarkRunOptsWithFn<T, T2>, "fn">,
+  fn?: (arg: T) => void | Promise<void>
+) {
+  const {fn: benchTask, ...opts} = coerceToOptsObj(idOrOpts, fn);
+  const currentSuite = getCurrentSuite();
+
+  const globalOptions = store.getGlobalOptions() ?? {};
+  const parentOptions = store.getOptions(getCurrentSuite()) ?? {};
+  const options = {...globalOptions, ...parentOptions, ...opts};
+  const {timeoutBench, maxMs, minMs} = options;
+
+  let timeout = timeoutBench ?? optionsDefault.timeoutBench;
+  if (maxMs && maxMs > timeout) {
+    timeout = maxMs * 1.5;
+  }
+
+  if (minMs && minMs > timeout) {
+    timeout = minMs * 1.5;
+  }
+
+  const task = currentSuite.task(opts.id, {
+    skip: opts.skip ?? this.skip,
+    only: opts.only ?? this.only,
+    sequential: true,
+    concurrent: false,
+    timeout,
+    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"));
+      }
+    },
+  });
+
+  store.setOptions(task, opts);
+});
+
+function createBenchmarkFunction(
+  fn: <T, T2>(
+    this: Record<"skip" | "only", boolean | undefined>,
+    idOrOpts: string | PartialBy<BenchmarkRunOptsWithFn<T, T2>, "fn">,
+    fn?: (arg: T) => void | Promise<void>
+  ) => void
+): BenchApi {
+  return createChainable(["skip", "only"], fn) as BenchApi;
+}
+
+function coerceToOptsObj<T, T2>(
+  idOrOpts: string | PartialBy<BenchmarkRunOptsWithFn<T, T2>, "fn">,
+  fn?: (arg: T) => void | Promise<void>
+): BenchmarkRunOptsWithFn<T, T2> {
+  let opts: BenchmarkRunOptsWithFn<T, T2>;
+
+  if (typeof idOrOpts === "string") {
+    if (!fn) throw Error("fn arg must be set");
+    opts = {id: idOrOpts, fn, threshold: optionsDefault.threshold};
+  } else {
+    if (fn) {
+      opts = {...idOrOpts, fn};
+    } else {
+      const optsWithFn = idOrOpts as BenchmarkRunOptsWithFn<T, T2>;
+      if (!optsWithFn.fn) throw Error("opts.fn arg must be set");
+      opts = optsWithFn;
+    }
+  }
+
+  return opts;
+}
+
+/**
+ * Customize benchmark opts for a describe block
+ * ```ts
+ * describe("suite A1", function () {
+ *   setBenchOpts({runs: 100});
+ *   // 100 runs
+ *   itBench("bench A1.1", function() {});
+ *   itBench("bench A1.2", function() {});
+ *   // 300 runs
+ *   itBench({id: "bench A1.3", runs: 300}, function() {});
+ *
+ *   // Supports nesting, child has priority over parent.
+ *   // Arrow functions can be used, won't break it.
+ *   describe("suite A2", () => {
+ *     setBenchOpts({runs: 200});
+ *     // 200 runs.
+ *     itBench("bench A2.1", () => {});
+ *   })
+ * })
+ * ```
+ */
+export function setBenchOpts(opts: BenchmarkOpts): void {
+  store.setOptions(getCurrentSuite(), opts);
+}
+
+export const setBenchmarkOptions = setBenchOpts;

src/benchmark/globalState.ts

@@ -0,0 +1,44 @@
+import type {Suite, SuiteCollector, Task} from "@vitest/runner";
+import {BenchmarkResult, BenchmarkOpts, BenchmarkResults} from "../types.js";
+
+/**t
+ * Map of results by root suite.
+ */
+const results = new Map<string, BenchmarkResult>();
+
+/**
+ * Global opts from CLI
+ */
+let globalOpts: BenchmarkOpts | undefined;
+
+/**
+ * Map to persist options set in describe blocks
+ */
+const optsMap = new WeakMap<object, BenchmarkOpts>();
+
+export const store = {
+  getResult(id: string): BenchmarkResult | undefined {
+    return results.get(id);
+  },
+  setResult(id: string, result: BenchmarkResult): void {
+    results.set(id, result);
+  },
+  getAllResults(): BenchmarkResults {
+    return [...results.values()];
+  },
+  getOptions(suite: Task | Suite | SuiteCollector): BenchmarkOpts | undefined {
+    return optsMap.get(suite);
+  },
+  setOptions(suite: Task | Suite | SuiteCollector, opts: BenchmarkOpts): void {
+    optsMap.set(suite, opts);
+  },
+  removeOptions(suite: Task | Suite): void {
+    optsMap.delete(suite);
+  },
+  setGlobalOptions(opts: Partial<BenchmarkOpts>): void {
+    globalOpts = opts;
+  },
+  getGlobalOptions(): BenchmarkOpts | undefined {
+    return globalOpts;
+  },
+};

src/benchmark/index.ts

@@ -0,0 +1 @@
+export * from "./benchmarkFn.js";

src/benchmark/reporter.ts

@@ -0,0 +1,107 @@
+import {Task, Suite, File} from "@vitest/runner";
+import {color, consoleLog, symbols} from "../utils/output.js";
+import {store} from "./globalState.js";
+import {Benchmark, BenchmarkOpts, BenchmarkResult} from "../types.js";
+import {formatResultRow} from "./format.js";
+import {optionsDefault} from "../cli/options.js";
+
+export class BenchmarkReporter {
+  indents = 0;
+  failed = 0;
+  passed = 0;
+  skipped = 0;
+
+  readonly prevResults: Map<string, BenchmarkResult>;
+  readonly threshold: number;
+
+  constructor({prevBench, benchmarkOpts}: {prevBench: Benchmark | null; benchmarkOpts: BenchmarkOpts}) {
+    this.prevResults = new Map<string, BenchmarkResult>();
+    this.threshold = benchmarkOpts.threshold ?? optionsDefault.threshold;
+
+    if (prevBench) {
+      for (const bench of prevBench.results) {
+        this.prevResults.set(bench.id, bench);
+      }
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  onTestStarted(_task: Task): void {
+    // this.log(task.name, "started");
+  }
+
+  onTestFinished(task: Task): void {
+    const {result} = task;
+
+    if (!result) {
+      consoleLog(`${this.indent()}${color("pending", " - %s")}`, `${task.name} - can not find result`);
+      return;
+    }
+
+    switch (result.state) {
+      case "skip": {
+        this.skipped++;
+        consoleLog(`${this.indent()}${color("pending", "  - %s")}`, task.name);
+        break;
+      }
+      case "fail": {
+        this.failed++;
+        consoleLog(this.indent() + color("fail", "  %d) %s"), ++this.failed, task.name);
+        consoleLog(task.result?.errors);
+        break;
+      }
+      case "pass": {
+        try {
+          const result = store.getResult(task.name);
+
+          if (result) {
+            // Render benchmark
+            const prevResult = this.prevResults.get(result.id) ?? null;
+
+            const resultRow = formatResultRow(result, prevResult, result.threshold ?? this.threshold);
+            const fmt = this.indent() + color("checkmark", "  " + symbols.ok) + " " + resultRow;
+            consoleLog(fmt);
+          } else {
+            // Render regular test
+            const fmt = this.indent() + color("checkmark", "  " + symbols.ok) + color("pass", " %s");
+            consoleLog(fmt, task.name);
+          }
+          this.passed++;
+        } catch (e) {
+          this.failed++;
+          consoleLog(e);
+          process.exitCode = 1;
+          throw e;
+        }
+      }
+    }
+  }
+
+  onSuiteStarted(suite: Suite): void {
+    this.indents++;
+    consoleLog(color("suite", "%s%s"), this.indent(), suite.name);
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  onSuiteFinished(_suite: Suite): void {
+    --this.indents;
+
+    if (this.indents === 1) {
+      consoleLog();
+    }
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  onComplete(_files: File[]): void {
+    consoleLog();
+    this.indents += 2;
+    consoleLog(color("checkmark", "%s%s"), this.indent(), `${this.passed} passing`);
+    consoleLog(color("fail", "%s%s"), this.indent(), `${this.failed} failed`);
+    consoleLog(color("pending", "%s%s"), this.indent(), `${this.skipped} pending`);
+    consoleLog();
+  }
+
+  protected indent(): string {
+    return Array(this.indents).join("  ");
+  }
+}