Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(REST): dynamic rate limit offsets #10099

Merged
merged 5 commits into from
Jan 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion packages/rest/__tests__/RequestManager.test.ts
Original file line number Diff line number Diff line change
@@ -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();
Expand Down Expand Up @@ -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);
});
1 change: 0 additions & 1 deletion packages/rest/src/lib/REST.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
5 changes: 3 additions & 2 deletions packages/rest/src/lib/handlers/BurstHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

/**
Expand Down Expand Up @@ -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) {
Expand Down
25 changes: 16 additions & 9 deletions packages/rest/src/lib/handlers/SequentialHandler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -104,8 +104,9 @@ export class SequentialHandler implements IHandler {
/**
* The time until queued requests can continue
*/
private get timeToReset(): number {
return this.reset + this.manager.options.offset - Date.now();
private getTimeToReset(routeId: RouteData): number {
const offset = normalizeRateLimitOffset(this.manager.options.offset, routeId.bucketRoute);
return this.reset + offset - Date.now();
}

/**
Expand Down Expand Up @@ -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
Expand All @@ -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);
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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, {
Expand Down
7 changes: 6 additions & 1 deletion packages/rest/src/lib/utils/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
18 changes: 17 additions & 1 deletion packages/rest/src/lib/utils/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
}
Loading