feat(REST): dynamic rate limit offsets (#10099)

* feat(REST): dynamic rate limit offsets

* chore: update tests

* chore: better doc comment

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

* fix: don't overlook globalReset

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

---------

Co-authored-by: Jiralite <33201955+Jiralite@users.noreply.github.com>
Co-authored-by: ckohen <chaikohen@gmail.com>
Co-authored-by: kodiakhq[bot] <49736102+kodiakhq[bot]@users.noreply.github.com>

packages/rest/__tests__/RequestManager.test.ts

@@ -1,6 +1,7 @@
 import { MockAgent, setGlobalDispatcher, type Interceptable } from 'undici';
 import { beforeEach, afterEach, test, expect } from 'vitest';
 import { REST } from '../src/index.js';
+import { normalizeRateLimitOffset } from '../src/lib/utils/utils.js';
 import { genPath } from './util.js';
 
 const api = new REST();
@@ -36,5 +37,5 @@ test('no token', async () => {
 test('negative offset', () => {
 	const badREST = new REST({ offset: -5_000 });
 
-	expect(badREST.options.offset).toEqual(0);
+	expect(normalizeRateLimitOffset(badREST.options.offset, 'hehe :3')).toEqual(0);
 });

packages/rest/src/lib/REST.ts

@@ -77,7 +77,6 @@ export class REST extends AsyncEventEmitter<RestEvents> {
 		super();
 		this.cdn = new CDN(options.cdn ?? DefaultRestOptions.cdn);
 		this.options = { ...DefaultRestOptions, ...options };
-		this.options.offset = Math.max(0, this.options.offset);
 		this.globalRemaining = Math.max(1, this.options.globalRequestsPerSecond);
 		this.agent = options.agent ?? null;
 

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

@@ -3,7 +3,7 @@ import type { REST } from '../REST.js';
 import type { IHandler } from '../interfaces/Handler.js';
 import { RESTEvents } from '../utils/constants.js';
 import type { ResponseLike, HandlerRequestData, RouteData, RateLimitData } from '../utils/types.js';
-import { onRateLimit, sleep } from '../utils/utils.js';
+import { normalizeRateLimitOffset, onRateLimit, sleep } from '../utils/utils.js';
 import { handleErrors, incrementInvalidCount, makeNetworkRequest } from './Shared.js';
 
 /**
@@ -90,7 +90,8 @@ export class BurstHandler implements IHandler {
 		const retry = res.headers.get('Retry-After');
 
 		// Amount of time in milliseconds until we should retry if rate limited (globally or otherwise)
-		if (retry) retryAfter = Number(retry) * 1_000 + this.manager.options.offset;
+		const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
+		if (retry) retryAfter = Number(retry) * 1_000 + offset;
 
 		// Count the invalid requests
 		if (status === 401 || status === 403 || status === 429) {

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

@@ -4,7 +4,7 @@ import type { REST } from '../REST.js';
 import type { IHandler } from '../interfaces/Handler.js';
 import { RESTEvents } from '../utils/constants.js';
 import type { RateLimitData, ResponseLike, HandlerRequestData, RouteData } from '../utils/types.js';
-import { hasSublimit, onRateLimit, sleep } from '../utils/utils.js';
+import { hasSublimit, normalizeRateLimitOffset, onRateLimit, sleep } from '../utils/utils.js';
 import { handleErrors, incrementInvalidCount, makeNetworkRequest } from './Shared.js';
 
 const enum QueueType {
@@ -104,8 +104,9 @@ export class SequentialHandler implements IHandler {
 	/**
 	 * The time until queued requests can continue
 	 */
-	private get timeToReset(): number {
+	private getTimeToReset(routeId: RouteData): number {
-		return this.reset + this.manager.options.offset - Date.now();
+		const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
+		return this.reset + offset - Date.now();
 	}
 
 	/**
@@ -209,9 +210,11 @@ export class SequentialHandler implements IHandler {
 			let delay: Promise<void>;
 
 			if (isGlobal) {
+				const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
+
 				// Set RateLimitData based on the global limit
 				limit = this.manager.options.globalRequestsPerSecond;
-				timeout = this.manager.globalReset + this.manager.options.offset - Date.now();
+				timeout = this.manager.globalReset + offset - Date.now();
 				// If this is the first task to reach the global timeout, set the global delay
 				if (!this.manager.globalDelay) {
 					// The global delay function clears the global delay state when it is resolved
@@ -222,7 +225,7 @@ export class SequentialHandler implements IHandler {
 			} else {
 				// Set RateLimitData based on the route-specific limit
 				limit = this.limit;
-				timeout = this.timeToReset;
+				timeout = this.getTimeToReset(routeId);
 				delay = sleep(timeout);
 			}
 
@@ -284,15 +287,17 @@ export class SequentialHandler implements IHandler {
 		const retry = res.headers.get('Retry-After');
 		const scope = (res.headers.get('X-RateLimit-Scope') ?? 'user') as RateLimitData['scope'];
 
+		const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
+
 		// Update the total number of requests that can be made before the rate limit resets
 		this.limit = limit ? Number(limit) : Number.POSITIVE_INFINITY;
 		// Update the number of remaining requests that can be made before the rate limit resets
 		this.remaining = remaining ? Number(remaining) : 1;
 		// Update the time when this rate limit resets (reset-after is in seconds)
-		this.reset = reset ? Number(reset) * 1_000 + Date.now() + this.manager.options.offset : Date.now();
+		this.reset = reset ? Number(reset) * 1_000 + Date.now() + offset : Date.now();
 
 		// Amount of time in milliseconds until we should retry if rate limited (globally or otherwise)
-		if (retry) retryAfter = Number(retry) * 1_000 + this.manager.options.offset;
+		if (retry) retryAfter = Number(retry) * 1_000 + offset;
 
 		// Handle buckets via the hash header retroactively
 		if (hash && hash !== this.hash) {
@@ -341,13 +346,15 @@ export class SequentialHandler implements IHandler {
 			let timeout: number;
 
 			if (isGlobal) {
+				const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
+
 				// Set RateLimitData based on the global limit
 				limit = this.manager.options.globalRequestsPerSecond;
-				timeout = this.manager.globalReset + this.manager.options.offset - Date.now();
+				timeout = this.manager.globalReset + offset - Date.now();
 			} else {
 				// Set RateLimitData based on the route-specific limit
 				limit = this.limit;
-				timeout = this.timeToReset;
+				timeout = this.getTimeToReset(routeId);
 			}
 
 			await onRateLimit(this.manager, {

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

@@ -90,7 +90,7 @@ export interface RESTOptions {
 	 *
 	 * @defaultValue `50`
 	 */
-	offset: number;
+	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
@@ -191,6 +191,11 @@ export interface RateLimitData {
  */
 export type RateLimitQueueFilter = (rateLimitData: RateLimitData) => Awaitable<boolean>;
 
+/**
+ * A function that determines the rate limit offset for a given request.
+ */
+export type GetRateLimitOffsetFunction = (route: string) => 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,8 @@ import type { RESTPatchAPIChannelJSONBody, Snowflake } from 'discord-api-types/v
 import type { REST } from '../REST.js';
 import { RateLimitError } from '../errors/RateLimitError.js';
 import { DEPRECATION_WARNING_PREFIX } from './constants.js';
-import { RequestMethod, type RateLimitData, type ResponseLike } from './types.js';
+import { RequestMethod } from './types.js';
+import type { GetRateLimitOffsetFunction, RateLimitData, ResponseLike } from './types.js';
 
 function serializeSearchParam(value: unknown): string | null {
 	switch (typeof value) {
@@ -156,3 +157,18 @@ export function deprecationWarning(message: string) {
 		process.emitWarning(message, DEPRECATION_WARNING_PREFIX);
 	}
 }
+
+/**
+ * Normalizes the offset for rate limits. Applies a Math.max(0, N) to prevent negative offsets,
+ * also deals with callbacks.
+ *
+ * @internal
+ */
+export function normalizeRateLimitOffset(offset: GetRateLimitOffsetFunction | number, route: string): number {
+	if (typeof offset === 'number') {
+		return Math.max(0, offset);
+	}
+
+	const result = offset(route);
+	return Math.max(0, result);
+}