fix: properly handle async function at all benchmark steps (#133)

Signed-off-by: Jérôme Benoit <[email protected]>
Co-authored-by: Jérôme Benoit <[email protected]>

src/task.ts

@@ -14,6 +14,7 @@ import type {
 import {
   absoluteDeviation,
   getVariance,
+  isAsyncFnResource,
   isAsyncTask,
   medianSorted,
   quantileSorted,
@@ -32,6 +33,9 @@ export default class Task extends EventTarget {
    */
   name: string;
 
+  /**
+   * Task function
+   */
   fn: Fn;
 
   /*
@@ -69,34 +73,52 @@ export default class Task extends EventTarget {
     const samples: number[] = [];
     if (this.opts.beforeAll != null) {
       try {
-        await this.opts.beforeAll.call(this);
+        if (await isAsyncFnResource(this.opts.beforeAll)) {
+          await this.opts.beforeAll.call(this);
+        } else {
+          this.opts.beforeAll.call(this);
+        }
       } catch (error) {
         return { error };
       }
     }
-    const isAsync = await isAsyncTask(this);
 
+    const asyncBeforeEach = this.opts.beforeEach != null
+      && (await isAsyncFnResource(this.opts.beforeEach));
+    const asyncTask = await isAsyncTask(this);
+    const asyncAfterEach = this.opts.afterEach != null
+      && (await isAsyncFnResource(this.opts.afterEach));
+
+    // TODO: factor out
     const executeTask = async () => {
       if (this.opts.beforeEach != null) {
-        await this.opts.beforeEach.call(this);
+        if (asyncBeforeEach) {
+          await this.opts.beforeEach.call(this);
+        } else {
+          this.opts.beforeEach.call(this);
+        }
       }
 
       let taskTime = 0;
-      if (isAsync) {
+      if (asyncTask) {
         const taskStart = this.bench.now();
-        await this.fn.call(this);
+        await this.fn();
         taskTime = this.bench.now() - taskStart;
       } else {
         const taskStart = this.bench.now();
-        this.fn.call(this);
+        this.fn();
         taskTime = this.bench.now() - taskStart;
       }
 
       samples.push(taskTime);
       totalTime += taskTime;
 
       if (this.opts.afterEach != null) {
-        await this.opts.afterEach.call(this);
+        if (asyncAfterEach) {
+          await this.opts.afterEach.call(this);
+        } else {
+          this.opts.afterEach.call(this);
+        }
       }
     };
 
@@ -123,7 +145,11 @@ export default class Task extends EventTarget {
 
     if (this.opts.afterAll != null) {
       try {
-        await this.opts.afterAll.call(this);
+        if (await isAsyncFnResource(this.opts.afterAll)) {
+          await this.opts.afterAll.call(this);
+        } else {
+          this.opts.afterAll.call(this);
+        }
       } catch (error) {
         return { error };
       }

src/utils.ts

@@ -7,58 +7,69 @@ export const hrtimeNow = () => nanoToMs(Number(process.hrtime.bigint()));
 
 export const now = () => performance.now();
 
-function isPromiseLike<T>(
+/**
+ * Checks if a value is a promise-like object.
+ *
+ * @param maybePromiseLike - the value to check
+ * @returns true if the value is a promise-like object
+ */
+const isPromiseLike = <T>(
   maybePromiseLike: any,
-): maybePromiseLike is PromiseLike<T> {
-  return (
-    maybePromiseLike !== null
-    && typeof maybePromiseLike === 'object'
-    && typeof maybePromiseLike.then === 'function'
-  );
-}
+): maybePromiseLike is PromiseLike<T> => maybePromiseLike !== null
+  && typeof maybePromiseLike === 'object'
+  && typeof maybePromiseLike.then === 'function';
 
-// eslint-disable-next-line @typescript-eslint/no-empty-function
-const AsyncFunctionConstructor = (async () => {}).constructor;
+type AsyncFunctionType<A extends unknown[], R> = (...args: A) => PromiseLike<R>;
 
 /**
- * An async function check method only consider runtime support async syntax
+ * An async function check helper only considering runtime support async syntax
+ *
+ * @param fn - the function to check
+ * @returns true if the function is an async function
  */
-export const isAsyncFunction = (fn: Fn) => fn.constructor === AsyncFunctionConstructor;
+const isAsyncFunction = (
+  fn: Fn,
+  // eslint-disable-next-line @typescript-eslint/no-empty-function
+): fn is AsyncFunctionType<unknown[], unknown> => fn?.constructor === (async () => {}).constructor;
 
-export const isAsyncTask = async (task: Task) => {
-  if (isAsyncFunction(task.fn)) {
+/**
+ * An async function check helper considering runtime support async syntax and promise return
+ *
+ * @param fn - the function to check
+ * @returns true if the function is an async function or returns a promise
+ */
+export const isAsyncFnResource = async (fn: Fn): Promise<boolean> => {
+  if (fn == null) {
+    return false;
+  }
+  if (isAsyncFunction(fn)) {
     return true;
   }
   try {
-    if (task.opts.beforeEach != null) {
-      try {
-        await task.opts.beforeEach.call(task);
-      } catch (e) {
-        // ignore
-      }
-    }
-    const call = task.fn();
-    const promiseLike = isPromiseLike(call);
+    const fnCall = fn();
+    const promiseLike = isPromiseLike(fnCall);
     if (promiseLike) {
+      // silence promise rejection
       try {
-        await call;
-      } catch (e) {
-        // ignore
-      }
-    }
-    if (task.opts.afterEach != null) {
-      try {
-        await task.opts.afterEach.call(task);
-      } catch (e) {
+        await fnCall;
+      } catch {
         // ignore
       }
     }
     return promiseLike;
-  } catch (e) {
+  } catch {
     return false;
   }
 };
 
+/**
+ * An async task check helper considering runtime support async syntax and promise return
+ *
+ * @param task - the task to check
+ * @returns true if the task is an async task
+ */
+export const isAsyncTask = async (task: Task): Promise<boolean> => isAsyncFnResource(task?.fn);
+
 /**
  * Computes the average of a sample.
  *
@@ -69,7 +80,6 @@ export const average = (samples: number[]) => {
   if (samples.length === 0) {
     throw new Error('samples must not be empty');
   }
-
   return samples.reduce((a, b) => a + b, 0) / samples.length || 0;
 };
 

test/index.test.ts

@@ -312,10 +312,14 @@ test('task beforeAll, afterAll, beforeEach, afterEach', async () => {
   await bench.warmup();
   await bench.run();
 
-  expect(beforeAll.mock.calls.length).toBe(2);
-  expect(afterAll.mock.calls.length).toBe(2);
-  expect(beforeEach.mock.calls.length).toBe(iterations * 2 /* warmup + run */);
-  expect(afterEach.mock.calls.length).toBe(iterations * 2);
+  expect(beforeAll.mock.calls.length).toBe(4 /* async check + warmup + run */);
+  expect(afterAll.mock.calls.length).toBe(4 /* async check + warmup + run */);
+  expect(beforeEach.mock.calls.length).toBe(
+    2 + iterations * 2 /* async check + warmup + run */,
+  );
+  expect(afterEach.mock.calls.length).toBe(
+    2 + iterations * 2 /* async check + warmup + run */,
+  );
   expect(beforeEach.mock.calls.length).toBe(afterEach.mock.calls.length);
 });
 

test/isAsyncTask.test.ts

@@ -4,6 +4,16 @@ import { isAsyncTask } from '../src/utils';
 
 const bench = new Bench();
 
+test('isAsyncTask undefined', () => {
+  // @ts-expect-error: testing with undefined
+  expect(isAsyncTask(undefined)).resolves.toBe(false);
+});
+
+test('isAsyncTask null', () => {
+  // @ts-expect-error: testing with null
+  expect(isAsyncTask(null)).resolves.toBe(false);
+});
+
 test('isAsyncTask sync', () => {
   const task = new Task(bench, 'foo', () => 1);
   expect(isAsyncTask(task)).resolves.toBe(false);
@@ -26,7 +36,7 @@ test('isAsyncTask promiseLike', () => {
   expect(isAsyncTask(task)).resolves.toBe(true);
 });
 
-test('isAsyncTask promise with catch', () => {
+test('isAsyncTask promise with error', () => {
   const task = new Task(bench, 'foo', () => Promise.reject(new Error('foo')));
   expect(isAsyncTask(task)).resolves.toBe(true);
 });