From 3c231ae81a52b66940ba495f35fd59a76c65e306 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?A=2E=20Rom=C3=A1n?= <kyradiscord@gmail.com>
Date: Sun, 25 Sep 2022 20:44:03 +0200
Subject: [PATCH] feat: add `AbortSignal` support (#8672)

* feat: add `AbortSignal` support

* fix: move the expect earlier

* fix: pass signal

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
---
 .../rest/__tests__/RequestHandler.test.ts     | 29 ++++++++++++++++++-
 packages/rest/src/lib/RequestManager.ts       |  7 ++++-
 .../src/lib/handlers/SequentialHandler.ts     | 19 ++++++++++--
 3 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/packages/rest/__tests__/RequestHandler.test.ts b/packages/rest/__tests__/RequestHandler.test.ts
index 6c846d6f1..d9c5746a5 100644
--- a/packages/rest/__tests__/RequestHandler.test.ts
+++ b/packages/rest/__tests__/RequestHandler.test.ts
@@ -1,7 +1,7 @@
 /* eslint-disable id-length */
 /* eslint-disable promise/prefer-await-to-then */
 import { performance } from 'node:perf_hooks';
-import { setInterval, clearInterval } from 'node:timers';
+import { setInterval, clearInterval, setTimeout } from 'node:timers';
 import { MockAgent, setGlobalDispatcher } from 'undici';
 import type { Interceptable, MockInterceptor } from 'undici/types/mock-interceptor';
 import { beforeEach, afterEach, test, expect, vitest } from 'vitest';
@@ -548,3 +548,30 @@ test('malformedRequest', async () => {
 
 	await expect(api.get('/malformedRequest')).rejects.toBeInstanceOf(DiscordAPIError);
 });
+
+test('abort', async () => {
+	mockPool
+		.intercept({
+			path: genPath('/abort'),
+			method: 'GET',
+		})
+		.reply(200, { message: 'Hello World' }, responseOptions)
+		.delay(100)
+		.times(3);
+
+	const controller = new AbortController();
+	const [aP2, bP2, cP2] = [
+		api.get('/abort', { signal: controller.signal }),
+		api.get('/abort', { signal: controller.signal }),
+		api.get('/abort', { signal: controller.signal }),
+	];
+
+	await expect(aP2).resolves.toStrictEqual({ message: 'Hello World' });
+	controller.abort();
+
+	// Abort mid-execution:
+	await expect(bP2).rejects.toThrowError('Request aborted');
+
+	// Abort scheduled:
+	await expect(cP2).rejects.toThrowError('Request aborted');
+});
diff --git a/packages/rest/src/lib/RequestManager.ts b/packages/rest/src/lib/RequestManager.ts
index 10a30bb5a..16624b4ba 100644
--- a/packages/rest/src/lib/RequestManager.ts
+++ b/packages/rest/src/lib/RequestManager.ts
@@ -93,6 +93,10 @@ export interface RequestData {
 	 * Reason to show in the audit logs
 	 */
 	reason?: string;
+	/**
+	 * The signal to abort the queue entry or the REST call, where applicable
+	 */
+	signal?: AbortSignal | undefined;
 	/**
 	 * If this request should be versioned
 	 *
@@ -133,7 +137,7 @@ export interface InternalRequest extends RequestData {
 	method: RequestMethod;
 }
 
-export type HandlerRequestData = Pick<InternalRequest, 'auth' | 'body' | 'files'>;
+export type HandlerRequestData = Pick<InternalRequest, 'auth' | 'body' | 'files' | 'signal'>;
 
 /**
  * Parsed route data for an endpoint
@@ -338,6 +342,7 @@ export class RequestManager extends EventEmitter {
 			body: request.body,
 			files: request.files,
 			auth: request.auth !== false,
+			signal: request.signal,
 		});
 	}
 
diff --git a/packages/rest/src/lib/handlers/SequentialHandler.ts b/packages/rest/src/lib/handlers/SequentialHandler.ts
index fa432f672..ef90ef43b 100644
--- a/packages/rest/src/lib/handlers/SequentialHandler.ts
+++ b/packages/rest/src/lib/handlers/SequentialHandler.ts
@@ -175,7 +175,7 @@ export class SequentialHandler implements IHandler {
 		}
 
 		// Wait for any previous requests to be completed before this one is run
-		await queue.wait();
+		await queue.wait({ signal: requestData.signal });
 		// This set handles retroactively sublimiting requests
 		if (queueType === QueueType.Standard) {
 			if (this.#sublimitedQueue && hasSublimit(routeId.bucketRoute, requestData.body, options.method)) {
@@ -293,8 +293,17 @@ export class SequentialHandler implements IHandler {
 
 		const controller = new AbortController();
 		const timeout = setTimeout(() => controller.abort(), this.manager.options.timeout).unref();
-		let res: Dispatcher.ResponseData;
+		if (requestData.signal) {
+			// The type polyfill is required because Node.js's types are incomplete.
+			const signal = requestData.signal as PolyFillAbortSignal;
+			// If the user signal was aborted, abort the controller, else abort the local signal.
+			// The reason why we don't re-use the user's signal, is because users may use the same signal for multiple
+			// requests, and we do not want to cause unexpected side-effects.
+			if (signal.aborted) controller.abort();
+			else signal.addEventListener('abort', () => controller.abort());
+		}
 
+		let res: Dispatcher.ResponseData;
 		try {
 			res = await request(url, { ...options, signal: controller.signal });
 		} catch (error: unknown) {
@@ -492,3 +501,9 @@ export class SequentialHandler implements IHandler {
 		}
 	}
 }
+
+interface PolyFillAbortSignal {
+	readonly aborted: boolean;
+	addEventListener(type: 'abort', listener: () => void): void;
+	removeEventListener(type: 'abort', listener: () => void): void;
+}