feat(rest): callbacks for timeout and retry backoff (#11067)

* feat(rest): callbacks for timeout and retry backoff

* test: add tests for callback utils

* test: fix typo

Co-authored-by: Qjuh <76154676+Qjuh@users.noreply.github.com>

* fix(retryBackoff): efficient math

* docs: minor tweaks

* docs: captalisation

---------

Co-authored-by: Qjuh <76154676+Qjuh@users.noreply.github.com>
Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com>

packages/rest/__tests__/BurstHandler.test.ts

@@ -16,6 +16,7 @@ const api = new REST();
 
 let mockAgent: MockAgent;
 let mockPool: Interceptable;
+let serverOutage = true;
 
 beforeEach(() => {
 	mockAgent = new MockAgent();
@@ -138,3 +139,57 @@ test('Handle unexpected 429', async () => {
 	expect(await unexpectedLimit).toStrictEqual({ test: true });
 	expect(performance.now()).toBeGreaterThanOrEqual(previous + 1_000);
 });
+
+test('server responding too slow', async () => {
+	const api2 = new REST({ timeout: 1 }).setToken('A-Very-Really-Real-Token');
+
+	api2.setAgent(mockAgent);
+
+	mockPool
+		.intercept({
+			path: callbackPath,
+			method: 'POST',
+		})
+		.reply(200, '')
+		.delay(100)
+		.times(10);
+
+	const promise = api2.post('/interactions/1234567890123456789/totallyarealtoken/callback', {
+		auth: false,
+		body: { type: 4, data: { content: 'Reply' } },
+	});
+
+	await expect(promise).rejects.toThrowError('aborted');
+}, 1_000);
+
+test('Handle temp server outage', async () => {
+	mockPool
+		.intercept({
+			path: callbackPath,
+			method: 'POST',
+		})
+		.reply(() => {
+			if (serverOutage) {
+				serverOutage = false;
+
+				return {
+					statusCode: 500,
+					data: '',
+				};
+			}
+
+			return {
+				statusCode: 200,
+				data: { test: true },
+				responseOptions,
+			};
+		})
+		.times(2);
+
+	expect(
+		await api.post('/interactions/1234567890123456789/totallyarealtoken/callback', {
+			auth: false,
+			body: { type: 4, data: { content: 'Reply' } },
+		}),
+	).toStrictEqual({ test: true });
+});

packages/rest/__tests__/utils.test.ts

@@ -1,5 +1,8 @@
+/* eslint-disable unicorn/consistent-function-scoping */
 import { describe, test, expect } from 'vitest';
+import type { GetRateLimitOffsetFunction, GetRetryBackoffFunction, GetTimeoutFunction } from '../src/index.js';
 import { makeURLSearchParams } from '../src/index.js';
+import { normalizeRateLimitOffset, normalizeRetryBackoff, normalizeTimeout } from '../src/lib/utils/utils.js';
 
 describe('makeURLSearchParams', () => {
 	test('GIVEN undefined THEN returns empty URLSearchParams', () => {
@@ -86,3 +89,94 @@ describe('makeURLSearchParams', () => {
 		});
 	});
 });
+
+describe('option normalization functions', () => {
+	describe('rate limit offset', () => {
+		const func: GetRateLimitOffsetFunction = (route) => {
+			if (route === '/negative') return -150;
+			if (route === '/high') return 150;
+			return 50;
+		};
+
+		test('offset as number', () => {
+			expect(normalizeRateLimitOffset(-150, '/negative')).toEqual(0);
+			expect(normalizeRateLimitOffset(150, '/high')).toEqual(150);
+			expect(normalizeRateLimitOffset(50, '/normal')).toEqual(50);
+		});
+
+		test('offset as function', () => {
+			expect(normalizeRateLimitOffset(func, '/negative')).toEqual(0);
+			expect(normalizeRateLimitOffset(func, '/high')).toEqual(150);
+			expect(normalizeRateLimitOffset(func, '/normal')).toEqual(50);
+		});
+	});
+
+	describe('retry backoff', () => {
+		const body = {
+			content: 'yo',
+		};
+		const func: GetRetryBackoffFunction = (_route, statusCode, retryCount) => {
+			if (statusCode === null) return 0;
+			if (statusCode === 502) return 50;
+			if (retryCount === 0) return 0;
+			if (retryCount === 1) return 150;
+			if (retryCount === 2) return 500;
+			return null;
+		};
+
+		test('retry backoff as number', () => {
+			expect(normalizeRetryBackoff(0, '/test', null, 0, body)).toEqual(0);
+			expect(normalizeRetryBackoff(0, '/test', null, 1, body)).toEqual(0);
+			expect(normalizeRetryBackoff(0, '/test', null, 2, body)).toEqual(0);
+			expect(normalizeRetryBackoff(50, '/test', null, 0, body)).toEqual(50);
+			expect(normalizeRetryBackoff(50, '/test', null, 1, body)).toEqual(100);
+			expect(normalizeRetryBackoff(50, '/test', null, 2, body)).toEqual(200);
+		});
+
+		test('retry backoff as function', () => {
+			expect(normalizeRetryBackoff(func, '/test', null, 0, body)).toEqual(0);
+			expect(normalizeRetryBackoff(func, '/test', 502, 0, body)).toEqual(50);
+			expect(normalizeRetryBackoff(func, '/test', 500, 0, body)).toEqual(0);
+			expect(normalizeRetryBackoff(func, '/test', 500, 1, body)).toEqual(150);
+			expect(normalizeRetryBackoff(func, '/test', 500, 2, body)).toEqual(500);
+			expect(normalizeRetryBackoff(func, '/test', 500, 3, body)).toEqual(null);
+		});
+	});
+
+	describe('timeout', () => {
+		const body1 = {
+			attachments: [{ id: 1 }],
+		};
+		const body2 = {
+			content: 'yo',
+		};
+		const func: GetTimeoutFunction = (route, body) => {
+			if (
+				typeof body === 'object' &&
+				body &&
+				'attachments' in body &&
+				Array.isArray(body.attachments) &&
+				body.attachments.length
+			) {
+				return 1_000;
+			}
+
+			if (route === '/negative') return -150;
+			if (route === '/high') return 150;
+			return 50;
+		};
+
+		test('timeout as number', () => {
+			expect(normalizeTimeout(-150, '/negative', body1)).toEqual(0);
+			expect(normalizeTimeout(150, '/high', body1)).toEqual(150);
+			expect(normalizeTimeout(50, '/normal', body1)).toEqual(50);
+		});
+
+		test('timeout as function', () => {
+			expect(normalizeTimeout(func, '/negative', body1)).toEqual(1_000);
+			expect(normalizeTimeout(func, '/negative', body2)).toEqual(0);
+			expect(normalizeTimeout(func, '/high', body2)).toEqual(150);
+			expect(normalizeTimeout(func, '/normal', body2)).toEqual(50);
+		});
+	});
+});

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

@@ -141,7 +141,7 @@ export class BurstHandler implements IHandler {
 			// Since this is not a server side issue, the next request should pass, so we don't bump the retries counter
 			return this.runRequest(routeId, url, options, requestData, retries);
 		} else {
-			const handled = await handleErrors(this.manager, res, method, url, requestData, retries);
+			const handled = await handleErrors(this.manager, res, method, url, requestData, retries, routeId);
 			if (handled === null) {
 				// eslint-disable-next-line no-param-reassign
 				return this.runRequest(routeId, url, options, requestData, ++retries);

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

@@ -420,7 +420,7 @@ export class SequentialHandler implements IHandler {
 			// Since this is not a server side issue, the next request should pass, so we don't bump the retries counter
 			return this.runRequest(routeId, url, options, requestData, retries);
 		} else {
-			const handled = await handleErrors(this.manager, res, method, url, requestData, retries);
+			const handled = await handleErrors(this.manager, res, method, url, requestData, retries, routeId);
 			if (handled === null) {
 				// eslint-disable-next-line no-param-reassign
 				return this.runRequest(routeId, url, options, requestData, ++retries);

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

@@ -5,7 +5,7 @@ import { DiscordAPIError } from '../errors/DiscordAPIError.js';
 import { HTTPError } from '../errors/HTTPError.js';
 import { RESTEvents } from '../utils/constants.js';
 import type { ResponseLike, HandlerRequestData, RouteData } from '../utils/types.js';
-import { parseResponse, shouldRetry } from '../utils/utils.js';
+import { normalizeRetryBackoff, normalizeTimeout, parseResponse, shouldRetry, sleep } from '../utils/utils.js';
 
 let authFalseWarningEmitted = false;
 
@@ -65,7 +65,10 @@ export async function makeNetworkRequest(
 	retries: number,
 ) {
 	const controller = new AbortController();
-	const timeout = setTimeout(() => controller.abort(), manager.options.timeout);
+	const timeout = setTimeout(
+		() => controller.abort(),
+		normalizeTimeout(manager.options.timeout, routeId.bucketRoute, requestData.body),
+	);
 	if (requestData.signal) {
 		// 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
@@ -81,6 +84,21 @@ export async function makeNetworkRequest(
 		if (!(error instanceof Error)) throw error;
 		// Retry the specified number of times if needed
 		if (shouldRetry(error) && retries !== manager.options.retries) {
+			const backoff = normalizeRetryBackoff(
+				manager.options.retryBackoff,
+				routeId.bucketRoute,
+				null,
+				retries,
+				requestData.body,
+			);
+			if (backoff === null) {
+				throw error;
+			}
+
+			if (backoff > 0) {
+				await sleep(backoff);
+			}
+
 			// Retry is handled by the handler upon receiving null
 			return null;
 		}
@@ -117,6 +135,7 @@ export async function makeNetworkRequest(
  * @param url - The fully resolved url to make the request to
  * @param requestData - Extra data from the user's request needed for errors and additional processing
  * @param retries - The number of retries this request has already attempted (recursion occurs on the handler)
+ * @param routeId - The generalized API route with literal ids for major parameters
  * @returns The response if the status code is not handled or null to request a retry
  */
 export async function handleErrors(
@@ -126,11 +145,27 @@ export async function handleErrors(
 	url: string,
 	requestData: HandlerRequestData,
 	retries: number,
+	routeId: RouteData,
 ) {
 	const status = res.status;
 	if (status >= 500 && status < 600) {
 		// Retry the specified number of times for possible server side issues
 		if (retries !== manager.options.retries) {
+			const backoff = normalizeRetryBackoff(
+				manager.options.retryBackoff,
+				routeId.bucketRoute,
+				status,
+				retries,
+				requestData.body,
+			);
+			if (backoff === null) {
+				throw new HTTPError(status, res.statusText, method, url, requestData);
+			}
+
+			if (backoff > 0) {
+				await sleep(backoff);
+			}
+
 			return null;
 		}
 

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

@@ -25,6 +25,7 @@ export const DefaultRestOptions = {
 	offset: 50,
 	rejectOnRateLimit: null,
 	retries: 3,
+	retryBackoff: 0,
 	timeout: 15_000,
 	userAgentAppendix: DefaultUserAgentAppendix,
 	version: APIVersion,

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

@@ -114,12 +114,18 @@ export interface RESTOptions {
 	 * @defaultValue `3`
 	 */
 	retries: number;
+	/**
+	 * The time to exponentially add before retrying a 5xx or aborted request
+	 *
+	 * @defaultValue `0`
+	 */
+	retryBackoff: GetRetryBackoffFunction | number;
 	/**
 	 * The time to wait in milliseconds before a request is aborted
 	 *
 	 * @defaultValue `15_000`
 	 */
-	timeout: number;
+	timeout: GetTimeoutFunction | number;
 	/**
 	 * Extra information to add to the user agent
 	 *
@@ -203,6 +209,30 @@ export type RateLimitQueueFilter = (rateLimitData: RateLimitData) => Awaitable<b
  */
 export type GetRateLimitOffsetFunction = (route: string) => number;
 
+/**
+ * A function that determines the backoff for a retry for a given request.
+ *
+ * @param route - The route that has encountered a server-side error
+ * @param statusCode - The status code received or `null` if aborted
+ * @param retryCount - The number of retries that have been attempted so far. The first call will be `0`
+ * @param requestBody - The body that was sent with the request
+ * @returns The delay for the current request or `null` to throw an error instead of retrying
+ */
+export type GetRetryBackoffFunction = (
+	route: string,
+	statusCode: number | null,
+	retryCount: number,
+	requestBody: unknown,
+) => number | null;
+
+/**
+ * A function that determines the timeout for a given request.
+ *
+ * @param route - The route that is being processed
+ * @param body - The body that will be sent with the request
+ */
+export type GetTimeoutFunction = (route: string, body: unknown) => number;
+
 export interface APIRequest {
 	/**
 	 * The data that was used to form the body of this request

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

@@ -2,7 +2,13 @@ import type { RESTPatchAPIChannelJSONBody, Snowflake } from 'discord-api-types/v
 import type { REST } from '../REST.js';
 import { RateLimitError } from '../errors/RateLimitError.js';
 import { RequestMethod } from './types.js';
-import type { GetRateLimitOffsetFunction, RateLimitData, ResponseLike } from './types.js';
+import type {
+	GetRateLimitOffsetFunction,
+	GetRetryBackoffFunction,
+	GetTimeoutFunction,
+	RateLimitData,
+	ResponseLike,
+} from './types.js';
 
 function serializeSearchParam(value: unknown): string | null {
 	// eslint-disable-next-line @typescript-eslint/switch-exhaustiveness-check
@@ -157,3 +163,39 @@ export function normalizeRateLimitOffset(offset: GetRateLimitOffsetFunction | nu
 	const result = offset(route);
 	return Math.max(0, result);
 }
+
+/**
+ * Normalizes the retry backoff used to add delay to retrying 5xx and aborted requests.
+ * Applies a Math.max(0, N) to prevent negative backoffs, also deals with callbacks.
+ *
+ * @internal
+ */
+export function normalizeRetryBackoff(
+	retryBackoff: GetRetryBackoffFunction | number,
+	route: string,
+	statusCode: number | null,
+	retryCount: number,
+	requestBody: unknown,
+): number | null {
+	if (typeof retryBackoff === 'number') {
+		return Math.max(0, retryBackoff) * (1 << retryCount);
+	}
+
+	// No need to Math.max as we'll only set the sleep timer if the value is > 0 (and not equal)
+	return retryBackoff(route, statusCode, retryCount, requestBody);
+}
+
+/**
+ * Normalizes the timeout for aborting requests. Applies a Math.max(0, N) to prevent negative timeouts,
+ * also deals with callbacks.
+ *
+ * @internal
+ */
+export function normalizeTimeout(timeout: GetTimeoutFunction | number, route: string, requestBody: unknown): number {
+	if (typeof timeout === 'number') {
+		return Math.max(0, timeout);
+	}
+
+	const result = timeout(route, requestBody);
+	return Math.max(0, result);
+}