Skip to content

Commit

Permalink
feat: add AbortSignal support (#8672)
Browse files Browse the repository at this point in the history
* feat: add `AbortSignal` support

* fix: move the expect earlier

* fix: pass signal

Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>
  • Loading branch information
kyranet and kodiakhq[bot] committed Sep 25, 2022
1 parent 9f63eb9 commit 3c231ae
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 4 deletions.
29 changes: 28 additions & 1 deletion 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';
Expand Down Expand Up @@ -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');
});
7 changes: 6 additions & 1 deletion packages/rest/src/lib/RequestManager.ts
Expand Up @@ -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
*
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -338,6 +342,7 @@ export class RequestManager extends EventEmitter {
body: request.body,
files: request.files,
auth: request.auth !== false,
signal: request.signal,
});
}

Expand Down
19 changes: 17 additions & 2 deletions packages/rest/src/lib/handlers/SequentialHandler.ts
Expand Up @@ -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)) {
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
}

0 comments on commit 3c231ae

Please sign in to comment.