From 21e98b0460820d34bf77e72f2585d0c692b69dea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Xalambr=C3=AD?= <hello@sergiodxa.com>
Date: Fri, 31 Mar 2023 16:04:21 -0500
Subject: [PATCH] Add timeout helper (#179)

---
 src/common/promise.ts       | 87 +++++++++++++++++++++++++++++++++++++
 test/common/promise.test.ts | 43 +++++++++++++++++-
 2 files changed, 129 insertions(+), 1 deletion(-)

diff --git a/src/common/promise.ts b/src/common/promise.ts
index 88ca552..b74ce74 100644
--- a/src/common/promise.ts
+++ b/src/common/promise.ts
@@ -44,3 +44,90 @@ export async function promiseHash<Hash extends PromiseHash>(
     )
   );
 }
+
+/**
+ * Used to uniquely identify a timeout
+ * @private
+ */
+let TIMEOUT = Symbol("TIMEOUT");
+
+/**
+ * Attach a timeout to any promise, if the timeout resolves first ignore the
+ * original promise and throw an error
+ * @param promise The promise to attach a timeout to
+ * @param options The options to use
+ * @param options.ms The number of milliseconds to wait before timing out
+ * @param options.controller An AbortController to abort the original promise
+ * @returns The result of the promise
+ * @throws TimeoutError If the timeout resolves first
+ * @example
+ * try {
+ *   let result = await timeout(
+ *     fetch("https://example.com"),
+ *     { ms: 100 }
+ *   );
+ * } catch (error) {
+ *   if (error instanceof TimeoutError) {
+ *     // Handle timeout
+ *   }
+ * }
+ * @example
+ * try {
+ *   let controller = new AbortController();
+ *   let result = await timeout(
+ *     fetch("https://example.com", { signal: controller.signal }),
+ *     { ms: 100, controller }
+ *   );
+ * } catch (error) {
+ *   if (error instanceof TimeoutError) {
+ *     // Handle timeout
+ *   }
+ * }
+ */
+export function timeout<Value>(
+  promise: Promise<Value>,
+  options: { controller?: AbortController; ms: number }
+): Promise<Value> {
+  return new Promise(async (resolve, reject) => {
+    let timer: NodeJS.Timeout | null = null;
+
+    try {
+      let result = await Promise.race([
+        promise,
+        new Promise((resolve) => {
+          timer = setTimeout(() => resolve(TIMEOUT), options.ms);
+        }),
+      ]);
+
+      if (timer) clearTimeout(timer);
+
+      if (result === TIMEOUT) {
+        if (options.controller) options.controller.abort();
+        return reject(new TimeoutError(`Timed out after ${options.ms}ms`));
+      }
+
+      return resolve(result as Awaited<Value>);
+    } catch (error) {
+      if (timer) clearTimeout(timer);
+      reject(error);
+    }
+  });
+}
+
+/**
+ * An error thrown when a timeout occurs
+ * @example
+ * try {
+ *   let result = await timeout(fetch("https://example.com"), { ms: 100 });
+ * } catch (error) {
+ *   if (error instanceof TimeoutError) {
+ *    // Handle timeout
+ *   }
+ * }
+ */
+export class TimeoutError extends Error {
+  constructor(message: string) {
+    super(message);
+    this.name = "TimeoutError";
+  }
+}
diff --git a/test/common/promise.test.ts b/test/common/promise.test.ts
index bd517fa..e1ac9ed 100644
--- a/test/common/promise.test.ts
+++ b/test/common/promise.test.ts
@@ -1,4 +1,4 @@
-import { promiseHash } from "../../src";
+import { promiseHash, timeout } from "../../src";
 
 describe(promiseHash, () => {
   test("should await all promises in a hash and return them with the same name", async () => {
@@ -23,3 +23,44 @@ describe(promiseHash, () => {
     });
   });
 });
+
+describe(timeout, () => {
+  test("resolves if the promise resolves before the timeout", async () => {
+    await expect(timeout(Promise.resolve(1), { ms: 1000 })).resolves.toBe(1);
+  });
+
+  test("rejects if the promise rejects before the timeout", async () => {
+    await expect(timeout(Promise.reject(1), { ms: 1000 })).rejects.toBe(1);
+  });
+
+  test("rejects if the timeout resolves first", async () => {
+    let timer: NodeJS.Timeout | null = null;
+
+    let promise = new Promise((resolve) => {
+      timer = setTimeout(resolve, 1000);
+    });
+
+    await expect(timeout(promise, { ms: 1 })).rejects.toThrow(
+      "Timed out after 1ms"
+    );
+
+    if (timer) clearTimeout(timer);
+  });
+
+  test("timeout aborts the controller", async () => {
+    let controller = new AbortController();
+    let timer: NodeJS.Timeout | null = null;
+
+    let promise = new Promise((resolve) => {
+      timer = setTimeout(resolve, 1000);
+    });
+
+    await expect(timeout(promise, { ms: 10, controller })).rejects.toThrow(
+      "Timed out after 10ms"
+    );
+
+    if (timer) clearTimeout(timer);
+
+    expect(controller.signal.aborted).toBe(true);
+  });
+});