From 98396ecbf7a71ab57678e0b9e09d31190dee1881 Mon Sep 17 00:00:00 2001
From: Daniil Bezuglov <81681662+picunada@users.noreply.github.com>
Date: Tue, 22 Aug 2023 20:10:06 +0700
Subject: [PATCH] feat: support `timeout` and `AbortController` (#268)

Co-authored-by: Pooya Parsa <pooya@pi0.io>
---
 README.md          | 10 ++++++++++
 src/fetch.ts       | 25 +++++++++++++++++++++----
 src/index.ts       |  3 ++-
 src/node.ts        |  8 ++++++--
 test/index.test.ts | 20 ++++++++++++++++++++
 5 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 3eb932c8..8d6fe405 100644
--- a/README.md
+++ b/README.md
@@ -128,6 +128,16 @@ await ofetch("http://google.com/404", {
 });
 ```
 
+## ✔️ Timeout
+
+You can specify `timeout` in milliseconds to automatically abort request after a timeout (default is disabled).
+
+```ts
+await ofetch("http://google.com/404", {
+  timeout: 3000, // Timeout after 3 seconds
+});
+```
+
 ## ✔️ Type Friendly
 
 Response can be type assisted:
diff --git a/src/fetch.ts b/src/fetch.ts
index ecbf46d3..882794ce 100644
--- a/src/fetch.ts
+++ b/src/fetch.ts
@@ -16,6 +16,7 @@ export interface CreateFetchOptions {
   defaults?: FetchOptions;
   fetch?: Fetch;
   Headers?: typeof Headers;
+  AbortController?: typeof AbortController;
 }
 
 export type FetchRequest = RequestInfo;
@@ -45,7 +46,8 @@ export interface FetchOptions<R extends ResponseType = ResponseType>
   responseType?: R;
   response?: boolean;
   retry?: number | false;
-
+  /** timeout in milliseconds */
+  timeout?: number;
   /** Delay between retries in milliseconds. */
   retryDelay?: number;
 
@@ -87,15 +89,21 @@ const retryStatusCodes = new Set([
 ]);
 
 export function createFetch(globalOptions: CreateFetchOptions = {}): $Fetch {
-  const { fetch = globalThis.fetch, Headers = globalThis.Headers } =
-    globalOptions;
+  const {
+    fetch = globalThis.fetch,
+    Headers = globalThis.Headers,
+    AbortController = globalThis.AbortController,
+  } = globalOptions;
 
   async function onError(context: FetchContext): Promise<FetchResponse<any>> {
     // Is Abort
     // If it is an active abort, it will not retry automatically.
     // https://developer.mozilla.org/en-US/docs/Web/API/DOMException#error_names
     const isAbort =
-      (context.error && context.error.name === "AbortError") || false;
+      (context.error &&
+        context.error.name === "AbortError" &&
+        !context.options.timeout) ||
+      false;
     // Retry
     if (context.options.retry !== false && !isAbort) {
       let retries;
@@ -111,9 +119,11 @@ export function createFetch(globalOptions: CreateFetchOptions = {}): $Fetch {
         if (retryDelay > 0) {
           await new Promise((resolve) => setTimeout(resolve, retryDelay));
         }
+        // Timeout
         return $fetchRaw(context.request, {
           ...context.options,
           retry: retries - 1,
+          timeout: context.options.timeout,
         });
       }
     }
@@ -184,6 +194,13 @@ export function createFetch(globalOptions: CreateFetchOptions = {}): $Fetch {
       }
     }
 
+    // TODO: Can we merge signals?
+    if (!context.options.signal && context.options.timeout) {
+      const controller = new AbortController();
+      setTimeout(() => controller.abort(), context.options.timeout);
+      context.options.signal = controller.signal;
+    }
+
     try {
       context.response = await fetch(
         context.request,
diff --git a/src/index.ts b/src/index.ts
index 8d68d839..4893fbf1 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -24,6 +24,7 @@ export const fetch =
   (() => Promise.reject(new Error("[ofetch] global.fetch is not supported!")));
 
 export const Headers = _globalThis.Headers;
+export const AbortController = _globalThis.AbortController;
 
-export const ofetch = createFetch({ fetch, Headers });
+export const ofetch = createFetch({ fetch, Headers, AbortController });
 export const $fetch = ofetch;
diff --git a/src/node.ts b/src/node.ts
index 0b2d45d2..9aa4a9dd 100644
--- a/src/node.ts
+++ b/src/node.ts
@@ -1,6 +1,9 @@
 import http from "node:http";
 import https, { AgentOptions } from "node:https";
-import nodeFetch, { Headers as _Headers } from "node-fetch-native";
+import nodeFetch, {
+  Headers as _Headers,
+  AbortController as _AbortController,
+} from "node-fetch-native";
 
 import { createFetch } from "./base";
 
@@ -33,6 +36,7 @@ export function createNodeFetch() {
 export const fetch = globalThis.fetch || createNodeFetch();
 
 export const Headers = globalThis.Headers || _Headers;
+export const AbortController = globalThis.AbortController || _AbortController;
 
-export const ofetch = createFetch({ fetch, Headers });
+export const ofetch = createFetch({ fetch, Headers, AbortController });
 export const $fetch = ofetch;
diff --git a/test/index.test.ts b/test/index.test.ts
index 5eceb631..680dca5f 100644
--- a/test/index.test.ts
+++ b/test/index.test.ts
@@ -65,6 +65,16 @@ describe("ofetch", () => {
       .use(
         "/408",
         eventHandler(() => createError({ status: 408 }))
+      )
+      .use(
+        "/timeout",
+        eventHandler(async () => {
+          await new Promise((resolve) => {
+            setTimeout(() => {
+              resolve(createError({ status: 408 }));
+            }, 1000 * 5);
+          });
+        })
       );
 
     listener = await listen(toNodeListener(app));
@@ -229,6 +239,16 @@ describe("ofetch", () => {
     expect(abortHandle()).rejects.toThrow(/aborted/);
   });
 
+  it("aborting on timeout", async () => {
+    const noTimeout = $fetch(getURL("timeout")).catch(() => "no timeout");
+    const timeout = $fetch(getURL("timeout"), {
+      timeout: 100,
+      retry: 0,
+    }).catch(() => "timeout");
+    const race = await Promise.race([noTimeout, timeout]);
+    expect(race).to.equal("timeout");
+  });
+
   it("deep merges defaultOptions", async () => {
     const _customFetch = $fetch.create({
       query: {