Add support for custom options

src/benchmark/benchmarkFn.ts

@@ -4,7 +4,8 @@ import {getCurrentSuite, Suite, SuiteCollector} from "@vitest/runner";
 import {createChainable} from "@vitest/runner/utils";
 import {store} from "./globalState.js";
 import {BenchmarkOpts} from "../types.js";
-import {runBenchFn} from "./runBenchFn.js";
+import {runBenchFn} from "./runBenchmarkFn.js";
+import {optionsDefault} from "../cli/options.js";
 
 export type BenchmarkRunOptsWithFn<T, T2> = BenchmarkOpts & {
   id: string;
@@ -37,47 +38,42 @@ export const bench = createBenchmarkFunction(function <T, T2>(
   fn?: (arg: T) => void | Promise<void>
 ) {
   const {fn: benchTask, ...opts} = coerceToOptsObj(idOrOpts, fn);
+  const currentSuite = getCurrentSuite();
 
-  const task = getCurrentSuite().task(opts.id, {
+  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: {
       "dapplion/benchmark": true,
     },
-    async handler(context) {
-      const parentSuite = context.task.suite;
-      const parentOpts = parentSuite ? store.getOptions(parentSuite) : {};
-
-      // TODO: Find better way to point to root suite
-      const rootSuite = context.task.suite;
-      const rootOpts = rootSuite ? store.getRootOptions(rootSuite) : {};
-
-      const fullOptions = Object.assign({}, rootOpts, parentOpts, opts);
-
+    async handler() {
       // Ensure bench id is unique
       if (store.getResult(opts.id) && !opts.skip) {
         throw Error(`test titles must be unique, duplicated: '${opts.id}'`);
       }
 
-      // Extend timeout if maxMs is set
-      // if (opts.timeoutBench !== undefined) {
-      // this.timeout(opts.timeoutBench);
-      // } else {
-      // const timeout = this.timeout();
-      // if (opts.maxMs && opts.maxMs > timeout) {
-      // this.timeout(opts.maxMs * 1.5);
-      // } else if (opts.minMs && opts.minMs > timeout) {
-      // this.timeout(opts.minMs * 1.5);
-      // }
-      // }
-
       // 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({...fullOptions, fn: benchTask}, persistRunsNs);
+      const {result, runsNs} = await runBenchFn({...options, fn: benchTask}, persistRunsNs);
 
       // Store result for:
       // - to persist benchmark data latter
@@ -130,7 +126,7 @@ function coerceToOptsObj<T, T2>(
 
   if (typeof idOrOpts === "string") {
     if (!fn) throw Error("fn arg must be set");
-    opts = {id: idOrOpts, fn};
+    opts = {id: idOrOpts, fn, threshold: optionsDefault.threshold};
   } else {
     if (fn) {
       opts = {...idOrOpts, fn};
@@ -143,3 +139,30 @@ function coerceToOptsObj<T, T2>(
 
   return opts;
 }
+
+/**
+ * Customize benchmark opts for a describe block. Affects only tests within that Mocha.Suite
+ * ```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

@@ -1,4 +1,4 @@
-import type {Suite, Task} from "@vitest/runner";
+import type {Suite, SuiteCollector, Task} from "@vitest/runner";
 import {BenchmarkResult, BenchmarkOpts, BenchmarkResults} from "../types.js";
 
 /**t
@@ -10,7 +10,7 @@ const results = new Map<string, BenchmarkResult>();
 /**
  * Global opts from CLI
  */
-const optsByRootSuite = new WeakMap<object, BenchmarkOpts>();
+let globalOpts: BenchmarkOpts | undefined;
 
 /**
  * Map to persist options set in describe blocks
@@ -27,16 +27,19 @@ export const store = {
   getAllResults(): BenchmarkResults {
     return [...results.values()];
   },
-  getOptions(suite: Task): BenchmarkOpts | undefined {
+  getOptions(suite: Task | Suite | SuiteCollector): BenchmarkOpts | undefined {
     return optsMap.get(suite);
   },
-  setOptions(suite: Task, opts: BenchmarkOpts): void {
+  setOptions(suite: Task | Suite | SuiteCollector, opts: BenchmarkOpts): void {
     optsMap.set(suite, opts);
   },
-  getRootOptions(suite: Suite): BenchmarkOpts | undefined {
-    return optsByRootSuite.get(suite);
+  removeOptions(suite: Task | Suite): void {
+    optsMap.delete(suite);
   },
-  setRootOptions(suite: Suite, opts: BenchmarkOpts): void {
-    optsByRootSuite.set(suite, opts);
+  setGlobalOptions(opts: Partial<BenchmarkOpts>): void {
+    globalOpts = opts;
+  },
+  getGlobalOptions(): BenchmarkOpts | undefined {
+    return globalOpts;
   },
 };

src/benchmark/reporter.ts

@@ -1,8 +1,9 @@
 import {Task, Suite, File} from "@vitest/runner";
 import {color, consoleLog, symbols} from "../utils/output.js";
 import {store} from "./globalState.js";
-import {Benchmark, BenchmarkResult} from "../types.js";
+import {Benchmark, BenchmarkOpts, BenchmarkResult} from "../types.js";
 import {formatResultRow} from "./format.js";
+import {optionsDefault} from "../cli/options.js";
 
 export class BenchmarkReporter {
   indents = 0;
@@ -13,9 +14,9 @@ export class BenchmarkReporter {
   readonly prevResults: Map<string, BenchmarkResult>;
   readonly threshold: number;
 
-  constructor(prevBench: Benchmark | null, threshold: number) {
+  constructor({prevBench, benchmarkOpts}: {prevBench: Benchmark | null; benchmarkOpts: BenchmarkOpts}) {
     this.prevResults = new Map<string, BenchmarkResult>();
-    this.threshold = threshold;
+    this.threshold = benchmarkOpts.threshold ?? optionsDefault.threshold;
 
     if (prevBench) {
       for (const bench of prevBench.results) {

src/benchmark/runner.ts

@@ -3,20 +3,21 @@ import {
   startTests,
   Suite,
   Task,
-  TaskResultPack,
   VitestRunner,
   VitestRunnerConfig,
   VitestRunnerImportSource,
 } from "@vitest/runner";
-import {Benchmark, BenchmarkResults} from "../types.js";
+import {Benchmark, BenchmarkOpts, BenchmarkResults} from "../types.js";
 import {BenchmarkReporter} from "./reporter.js";
 import {store} from "./globalState.js";
 
 export class BenchmarkRunner implements VitestRunner {
-  config: VitestRunnerConfig;
-  reporter: BenchmarkReporter;
+  readonly config: VitestRunnerConfig;
+  readonly reporter: BenchmarkReporter;
+  readonly prevBench: Benchmark | null;
+  readonly benchmarkOpts: BenchmarkOpts;
 
-  constructor(protected opts: {prevBench: Benchmark | null}) {
+  constructor({prevBench, benchmarkOpts}: {prevBench: Benchmark | null; benchmarkOpts: BenchmarkOpts}) {
     this.config = {
       root: "",
       sequence: {seed: 1234, hooks: "list", setupFiles: "list"},
@@ -27,7 +28,9 @@ export class BenchmarkRunner implements VitestRunner {
       setupFiles: [],
       retry: 0,
     };
-    this.reporter = new BenchmarkReporter(opts.prevBench, 0.2);
+    this.prevBench = prevBench;
+    this.benchmarkOpts = benchmarkOpts;
+    this.reporter = new BenchmarkReporter({prevBench, benchmarkOpts});
   }
 
   onBeforeRunSuite(suite: Suite): void {
@@ -57,12 +60,22 @@ export class BenchmarkRunner implements VitestRunner {
   }
 
   async process(files: string[]): Promise<BenchmarkResults> {
+    store.setGlobalOptions(this.benchmarkOpts);
+
     const res = await startTests(files, this);
 
-    if (res[0].result?.state === "pass") {
+    const passed = res.filter((r) => r.result?.state == "pass");
+    const skipped = res.filter((r) => r.result?.state == "skip");
+    const failed = res.filter((r) => r.result?.state == "fail");
+
+    if (failed.length > 0) {
+      throw failed[0].result?.errors;
+    }
+
+    if (passed.length + skipped.length === res.length) {
       return store.getAllResults();
     }
 
-    return store.getAllResults();
+    throw new Error("Some tests cause returned with unknown state");
   }
 }

src/cli/options.ts

@@ -3,6 +3,7 @@ import {StorageOptions, BenchmarkOpts, FileCollectionOptions} from "../types.js"
 
 export const optionsDefault = {
   threshold: 2,
+  timeoutBench: 10_000,
   historyLocalPath: "./benchmark_data",
   historyCacheKey: "benchmark_data",
 };
@@ -174,7 +175,8 @@ export const benchmarkOptions: ICliCommandOptions<CLIBenchmarkOptions> = {
   },
   timeoutBench: {
     type: "number",
-    description: "Hard timeout, enforced by mocha.",
+    description: "Hard timeout for each benchmark",
+    default: optionsDefault.timeoutBench,
     group: benchmarkGroup,
   },
 };

src/cli/run.ts

@@ -10,8 +10,7 @@ import {postGaComment} from "../github/comment.js";
 import {isGaRun} from "../github/context.js";
 import {BenchmarkRunner} from "../benchmark/runner.js";
 import {optionsDefault} from "./options.js";
-
-/* eslint-disable no-console */
+import {consoleLog} from "../utils/output.js";
 
 export async function run(opts_: FileCollectionOptions & StorageOptions & BenchmarkOpts): Promise<void> {
   const opts = Object.assign({}, optionsDefault, opts_);
@@ -21,33 +20,33 @@ export async function run(opts_: FileCollectionOptions & StorageOptions & Benchm
 
   // Retrieve history
   const historyProvider = getHistoryProvider(opts);
-  console.log(`Connected to historyProvider: ${historyProvider.providerInfo()}`);
+  consoleLog(`Connected to historyProvider: ${historyProvider.providerInfo()}`);
 
   // Select prev benchmark to compare against
   const compareWith = await resolveCompareWith(opts);
   const prevBench = await resolvePrevBenchmark(compareWith, historyProvider);
   if (prevBench) {
-    console.log(`Found previous benchmark for ${renderCompareWith(compareWith)}, at commit ${prevBench.commitSha}`);
+    consoleLog(`Found previous benchmark for ${renderCompareWith(compareWith)}, at commit ${prevBench.commitSha}`);
     validateBenchmark(prevBench);
   } else {
-    console.log(`No previous benchmark found for ${renderCompareWith(compareWith)}`);
+    consoleLog(`No previous benchmark found for ${renderCompareWith(compareWith)}`);
   }
 
   const {files, unmatchedFiles} = await collectFiles(opts).catch((err) => {
-    console.log("Error loading up spec patterns");
+    consoleLog("Error loading up spec patterns");
     throw err;
   });
 
   if (unmatchedFiles.length > 0) {
-    console.log(`Found unmatched files: \n${unmatchedFiles.join("\n")}\n`);
+    consoleLog(`Found unmatched files: \n${unmatchedFiles.join("\n")}\n`);
   }
 
   if (files.length === 0) {
-    console.log(`Can not find any matching spec file for ${opts.spec.join(",")}\n`);
+    consoleLog(`Can not find any matching spec file for ${opts.spec.join(",")}\n`);
     process.exit(1);
   }
 
-  const runner = new BenchmarkRunner({prevBench});
+  const runner = new BenchmarkRunner({prevBench, benchmarkOpts: opts});
   const results = await runner.process(files);
 
   if (results.length === 0) {
@@ -66,7 +65,7 @@ export async function run(opts_: FileCollectionOptions & StorageOptions & Benchm
   if (shouldPersist === true) {
     const refStr = github.context.ref || (await shell("git symbolic-ref HEAD"));
     const branch = parseBranchFromRef(refStr);
-    console.log(`Persisting new benchmark data for branch '${branch}' commit '${currBench.commitSha}'`);
+    consoleLog(`Persisting new benchmark data for branch '${branch}' commit '${currBench.commitSha}'`);
     // TODO: prune and limit total entries
     // appendBenchmarkToHistoryAndPrune(history, currBench, branch, opts);
     await historyProvider.writeLatestInBranch(branch, currBench);

src/index.ts

@@ -1,5 +1,5 @@
 import {suite, test} from "@vitest/runner";
 
-export {bench, itBench} from "./benchmark/index.js";
+export {bench, itBench, setBenchOpts, setBenchmarkOptions} from "./benchmark/index.js";
 export const describe = suite;
 export const it = test;

test/perf/iteration.test.ts

@@ -1,5 +1,5 @@
 import assert from "node:assert";
-import {itBench, describe, it} from "../../src/index.js";
+import {itBench, describe, it, setBenchOpts} from "../../src/index.js";
 
 // As of Jun 17 2021
 // Compare state root
@@ -9,7 +9,7 @@ import {itBench, describe, it} from "../../src/index.js";
 // byteArrayEquals with valueOf()                                         853971.0 ops/s      1.171000 us/op 9963051 runs    16.07 s
 
 describe("Array iteration", () => {
-  // setBenchOpts({maxMs: 60 * 1000, convergeFactor: 0.1 / 100});
+  setBenchOpts({maxMs: 60 * 1000, convergeFactor: 0.1 / 100});
 
   it("Regular test", () => {
     assert.strictEqual(1 + 2, 3);