feat(rest)!: add rejectOnRateLimit on a per-request basis (#11600)

BREAKING CHANGE: `RESTOptions#rejectOnRateLimit` now only accepts `boolean` or `RateLimitQueueFilter`

Co-authored-by: ckohen <chaikohen@gmail.com>

apps/proxy-container/src/index.ts

@@ -6,7 +6,7 @@ import { REST } from '@discordjs/rest';
 process.on('SIGINT', () => process.exit(0));
 
 // We want to let upstream handle retrying
-const api = new REST({ rejectOnRateLimit: () => true, retries: 0 });
+const api = new REST({ rejectOnRateLimit: true, retries: 0 });
 const server = createServer(proxyRequests(api));
 
 const port = Number.parseInt(process.env.PORT ?? '8080', 10);

packages/rest/__tests__/BurstHandler.test.ts

@@ -3,7 +3,7 @@
 import { MockAgent, setGlobalDispatcher } from 'undici';
 import type { Interceptable, MockInterceptor } from 'undici/types/mock-interceptor';
 import { beforeEach, afterEach, test, expect } from 'vitest';
-import { DiscordAPIError, REST, BurstHandlerMajorIdKey } from '../src/index.js';
+import { DiscordAPIError, RateLimitError, REST, BurstHandlerMajorIdKey } from '../src/index.js';
 import { BurstHandler } from '../src/lib/handlers/BurstHandler.js';
 import { genPath } from './util.js';
 
@@ -138,6 +138,33 @@ test('Handle unexpected 429', async () => {
 	expect(firstResolvedTime!).toBeGreaterThanOrEqual(previous + 1_000);
 });
 
+test('rejectOnRateLimit rejects on an unexpected 429', async () => {
+	mockPool
+		.intercept({
+			path: callbackPath,
+			method: 'POST',
+		})
+		.reply(() => ({
+			statusCode: 429,
+			data: '',
+			responseOptions: {
+				headers: {
+					'retry-after': '1',
+					via: '1.1 google',
+				},
+			},
+		}))
+		.times(1);
+
+	const promise = api.post('/interactions/1234567890123456789/totallyarealtoken/callback', {
+		auth: false,
+		body: { type: 4, data: { content: 'Reply' } },
+		rejectOnRateLimit: true,
+	});
+
+	await expect(promise).rejects.toBeInstanceOf(RateLimitError);
+});
+
 test('server responding too slow', async () => {
 	const api2 = new REST({ timeout: 1 }).setToken('A-Very-Really-Real-Token');
 

packages/rest/__tests__/RequestHandler.test.ts

@@ -12,7 +12,9 @@ let mockPool: Interceptable;
 
 const api = new REST({ timeout: 2_000, offset: 5 }).setToken('A-Very-Fake-Token');
 const invalidAuthApi = new REST({ timeout: 2_000 }).setToken('Definitely-Not-A-Fake-Token');
-const rateLimitErrorApi = new REST({ rejectOnRateLimit: ['/channels'] }).setToken('Obviously-Not-A-Fake-Token');
+const rateLimitErrorApi = new REST({
+	rejectOnRateLimit: (rateLimitData) => rateLimitData.route.startsWith('/channels'),
+}).setToken('Obviously-Not-A-Fake-Token');
 
 beforeEach(() => {
 	mockAgent = new MockAgent();
@@ -45,6 +47,7 @@ let sublimitHits = 0;
 let serverOutage = true;
 let unexpected429 = true;
 let unexpected429cf = true;
+let optOut429 = true;
 const sublimitIntervals: {
 	reset: NodeJS.Timeout | null;
 	retry: NodeJS.Timeout | null;
@@ -409,6 +412,111 @@ test('Handle unexpected 429 cloudflare', async () => {
 	expect(Date.now()).toBeGreaterThanOrEqual(previous + 1_000);
 });
 
+test('rejectOnRateLimit rejects on the pre-emptive throttle', async () => {
+	mockPool
+		.intercept({
+			path: genPath('/preemptive'),
+			method: 'GET',
+		})
+		.reply(() => ({
+			statusCode: 200,
+			data: { test: true },
+			responseOptions: {
+				headers: {
+					...responseOptions.headers,
+					'x-ratelimit-limit': '1',
+					'x-ratelimit-remaining': '0',
+					'x-ratelimit-reset-after': '0.5',
+					via: '1.1 google',
+				},
+			},
+		}))
+		.times(1);
+
+	expect(await api.get('/preemptive')).toStrictEqual({ test: true });
+	await expect(api.get('/preemptive', { rejectOnRateLimit: true })).rejects.toBeInstanceOf(RateLimitError);
+});
+
+test('rejectOnRateLimit rejects on an unexpected 429', async () => {
+	mockPool
+		.intercept({
+			path: genPath('/reject-429'),
+			method: 'GET',
+		})
+		.reply(() => ({
+			statusCode: 429,
+			data: '',
+			responseOptions: {
+				headers: {
+					'retry-after': '1',
+					'x-ratelimit-scope': 'shared',
+					via: '1.1 google',
+				},
+			},
+		}))
+		.times(1);
+
+	const rejectOnRateLimit = vitest.fn(() => true);
+
+	await expect(api.get('/reject-429', { rejectOnRateLimit })).rejects.toBeInstanceOf(RateLimitError);
+
+	expect(rejectOnRateLimit).toHaveBeenCalledTimes(1);
+	expect(rejectOnRateLimit).toHaveBeenCalledWith(
+		expect.objectContaining({
+			global: false,
+			method: 'GET',
+			route: '/reject-429',
+			majorParameter: 'global',
+			// 1_005 because of `offset: 5`
+			retryAfter: 1_005,
+			sublimitTimeout: 1_005,
+			scope: 'shared',
+		}),
+	);
+});
+
+test('Per-call rejectOnRateLimit takes precedence over the instance-wide one', async () => {
+	mockPool
+		.intercept({
+			path: genPath('/channels/1111111111111111111'),
+			method: 'GET',
+		})
+		.reply(() => ({
+			statusCode: 429,
+			data: '',
+			responseOptions: { headers: { 'retry-after': '1', via: '1.1 google' } },
+		}))
+		.times(1);
+
+	mockPool
+		.intercept({
+			path: genPath('/channels/2222222222222222222'),
+			method: 'GET',
+		})
+		.reply(() => {
+			if (optOut429) {
+				optOut429 = false;
+
+				return {
+					statusCode: 429,
+					data: '',
+					responseOptions: { headers: { 'retry-after': '1', via: '1.1 google' } },
+				};
+			}
+
+			return { statusCode: 200, data: { test: true }, responseOptions };
+		})
+		.times(2);
+
+	await expect(rateLimitErrorApi.get('/channels/1111111111111111111')).rejects.toBeInstanceOf(RateLimitError);
+
+	const previous = performance.now();
+	expect(await rateLimitErrorApi.get('/channels/2222222222222222222', { rejectOnRateLimit: false })).toStrictEqual({
+		test: true,
+	});
+	expect(performance.now()).toBeGreaterThanOrEqual(previous + 1_000);
+});
+
 test('Handle global rate limits', async () => {
 	mockPool
 		.intercept({

packages/rest/src/lib/REST.ts

@@ -264,6 +264,7 @@ export class REST extends AsyncEventEmitter<RestEvents> {
 			body: request.body,
 			files: request.files,
 			auth,
+			rejectOnRateLimit: request.rejectOnRateLimit,
 			signal: request.signal,
 		});
 	}

packages/rest/src/lib/handlers/BurstHandler.ts

@@ -105,19 +105,23 @@ export class BurstHandler implements IHandler {
 			const isGlobal = res.headers.has('X-RateLimit-Global');
 			const scope = (res.headers.get('X-RateLimit-Scope') ?? 'user') as RateLimitData['scope'];
 
-			await onRateLimit(this.manager, {
-				global: isGlobal,
-				method,
-				url,
-				route: routeId.bucketRoute,
-				majorParameter: this.majorParameter,
-				hash: this.hash,
-				limit: Number.POSITIVE_INFINITY,
-				timeToReset: retryAfter,
-				retryAfter,
-				sublimitTimeout: 0,
-				scope,
-			});
+			await onRateLimit(
+				this.manager,
+				{
+					global: isGlobal,
+					method,
+					url,
+					route: routeId.bucketRoute,
+					majorParameter: this.majorParameter,
+					hash: this.hash,
+					limit: Number.POSITIVE_INFINITY,
+					timeToReset: retryAfter,
+					retryAfter,
+					sublimitTimeout: 0,
+					scope,
+				},
+				requestData,
+			);
 
 			this.debug(
 				[

packages/rest/src/lib/handlers/SequentialHandler.ts

@@ -247,7 +247,7 @@ export class SequentialHandler implements IHandler {
 			// Let library users know they have hit a rate limit
 			this.manager.emit(RESTEvents.RateLimited, rateLimitData);
 			// Determine whether a RateLimitError should be thrown
-			await onRateLimit(this.manager, rateLimitData);
+			await onRateLimit(this.manager, rateLimitData, requestData);
 
 			// When not erroring, emit debug for what is happening
 			if (isGlobal) {
@@ -363,19 +363,23 @@ export class SequentialHandler implements IHandler {
 				timeout = this.getTimeToReset(routeId);
 			}
 
-			await onRateLimit(this.manager, {
-				global: isGlobal,
-				method,
-				url,
-				route: routeId.bucketRoute,
-				majorParameter: this.majorParameter,
-				hash: this.hash,
-				limit,
-				timeToReset: timeout,
-				retryAfter,
-				sublimitTimeout: sublimitTimeout ?? 0,
-				scope,
-			});
+			await onRateLimit(
+				this.manager,
+				{
+					global: isGlobal,
+					method,
+					url,
+					route: routeId.bucketRoute,
+					majorParameter: this.majorParameter,
+					hash: this.hash,
+					limit,
+					timeToReset: timeout,
+					retryAfter,
+					sublimitTimeout: sublimitTimeout ?? 0,
+					scope,
+				},
+				requestData,
+			);
 
 			this.debug(
 				[

packages/rest/src/lib/utils/constants.ts

@@ -23,7 +23,7 @@ export const DefaultRestOptions = {
 	invalidRequestWarningInterval: 0,
 	globalRequestsPerSecond: 50,
 	offset: 50,
-	rejectOnRateLimit: null,
+	rejectOnRateLimit: false,
 	retries: 3,
 	retryBackoff: 0,
 	timeout: 15_000,

packages/rest/src/lib/utils/types.ts

@@ -99,14 +99,16 @@ export interface RESTOptions {
 	 */
 	offset: GetRateLimitOffsetFunction | number;
 	/**
-	 * Determines how rate limiting and pre-emptive throttling should be handled.
-	 * When an array of strings, each element is treated as a prefix for the request route
-	 * (e.g. `/channels` to match any route starting with `/channels` such as `/channels/:id/messages`)
-	 * for which to throw {@link RateLimitError}s. All other request routes will be queued normally
+	 * The default policy determining how rate limiting and pre-emptive throttling should be handled.
 	 *
-	 * @defaultValue `null`
+	 * Pass `true` to throw a {@link RateLimitError} on every rate limit, `false` to wait every
+	 * rate limit out, or a filter to decide per rate limit.
+	 *
+	 * This can be overridden per request via the {@link RequestData.rejectOnRateLimit | rejectOnRateLimit} request option.
+	 *
+	 * @defaultValue `false`
 	 */
-	rejectOnRateLimit: RateLimitQueueFilter | string[] | null;
+	rejectOnRateLimit: RateLimitQueueFilter | boolean;
 	/**
 	 * The number of retries for errors with the 500 code, or errors
 	 * that timeout
@@ -339,6 +341,32 @@ export interface RequestData {
 	 * Reason to show in the audit logs
 	 */
 	reason?: string | undefined;
+	/**
+	 * Determines how a rate limit encountered while making this request should be handled.
+	 *
+	 * Pass `true` to throw a {@link RateLimitError} rather than wait, `false` to wait it out, or
+	 * a filter to decide based on rate limit data. Takes precedence over {@link RESTOptions.rejectOnRateLimit}, so
+	 * `false` opts this request out of an instance-wide policy. Leave it unset to inherit.
+	 *
+	 * @example
+	 * ```ts
+	 * // Fail rather than wait, no matter the rate limit
+	 * await rest.get(Routes.channel(channelId), { rejectOnRateLimit: true });
+	 *
+	 * // Give up rather than wait out a sublimit, which may be several minutes long
+	 * await rest.patch(Routes.channel(channelId), {
+	 * 	body: { name },
+	 * 	rejectOnRateLimit: (rateLimitData) => rateLimitData.sublimitTimeout > 0,
+	 * });
+	 *
+	 * // Spend at most 10 seconds waiting on rate limits
+	 * const deadline = Date.now() + 10_000;
+	 * await rest.get(Routes.channel(channelId), {
+	 * 	rejectOnRateLimit: (rateLimitData) => Date.now() + rateLimitData.retryAfter > deadline,
+	 * });
+	 * ```
+	 */
+	rejectOnRateLimit?: RateLimitQueueFilter | boolean | undefined;
 	/**
 	 * The signal to abort the queue entry or the REST call, where applicable
 	 */
@@ -381,7 +409,7 @@ export interface InternalRequest extends RequestData {
 	method: RequestMethod;
 }
 
-export interface HandlerRequestData extends Pick<InternalRequest, 'body' | 'files' | 'signal'> {
+export interface HandlerRequestData extends Pick<InternalRequest, 'body' | 'files' | 'rejectOnRateLimit' | 'signal'> {
 	auth: boolean | string;
 }
 

packages/rest/src/lib/utils/utils.ts

@@ -7,6 +7,7 @@ import type {
 	GetRateLimitOffsetFunction,
 	GetRetryBackoffFunction,
 	GetTimeoutFunction,
+	HandlerRequestData,
 	RateLimitData,
 	ResponseLike,
 } from './types.js';
@@ -150,15 +151,13 @@ export function shouldRetry(error: Error | NodeJS.ErrnoException) {
  *
  * @internal
  */
-export async function onRateLimit(manager: REST, rateLimitData: RateLimitData) {
-	const { options } = manager;
-	if (!options.rejectOnRateLimit) return;
+export async function onRateLimit(manager: REST, rateLimitData: RateLimitData, requestData: HandlerRequestData) {
+	// Explicit false opts out of `REST` level `rejectOnRateLimit`, only `undefined` falls back.
+	const policy = requestData.rejectOnRateLimit ?? manager.options.rejectOnRateLimit;
 
-	const shouldThrow =
-		typeof options.rejectOnRateLimit === 'function'
-			? await options.rejectOnRateLimit(rateLimitData)
-			: options.rejectOnRateLimit.some((route) => rateLimitData.route.startsWith(route.toLowerCase()));
-	if (shouldThrow) {
+	if (!policy) return;
+
+	if (policy === true || (await policy(rateLimitData))) {
 		throw new RateLimitError(rateLimitData);
 	}
 }