diff --git a/src/lib/rateLimit.ts b/src/lib/rateLimit.ts new file mode 100644 index 00000000..4e7e6543 --- /dev/null +++ b/src/lib/rateLimit.ts @@ -0,0 +1,179 @@ +import { RequestError, type CancelableRequest, type Response } from 'got'; +import { logger } from '../logger'; + +interface RateLimitInfo { + /** Total number of requests available per reset */ + limit: number | undefined; + /** Remaining number of requests until reset */ + remaining: number | undefined; + /** Time at which the limit will next reset */ + reset: Date | undefined; +} + +/** + * Wrap repeated `got` calls with rate limit handling. + */ +export class RateLimitClient { + private static readonly MAX_WAIT_TIME_MS = 3 * 60 * 1000; + + /** Total number of requests available per reset */ + public limit: number | undefined; + + /** Remaining number of requests until reset */ + public remaining: number | undefined; + + /** Time at which the limit will next reset */ + public reset: Date | undefined; + + /** Last response */ + protected lastResponse: Response | undefined; + + /** + * Create a new rate limit client. Best to create a new instance for each batch of calls to an endpoint. + */ + constructor() { + this.limit = undefined; + this.remaining = undefined; + this.reset = undefined; + this.lastResponse = undefined; + } + + /** + * Get the rate limit information from the headers + * + * @param headers - The headers from the got Response + * @returns The rate limit information + */ + static getRateLimitInfo(headers: Response['headers']): RateLimitInfo { + return { + limit: + typeof headers['x-ratelimit-limit'] === 'string' + ? Number.parseInt(headers['x-ratelimit-limit'], 10) + : undefined, + remaining: + typeof headers['x-ratelimit-remaining'] === 'string' + ? Number.parseInt(headers['x-ratelimit-remaining'], 10) + : undefined, + reset: + typeof headers['x-ratelimit-reset'] === 'string' + ? new Date(Number.parseInt(headers['x-ratelimit-reset'], 10) * 1000) + : undefined, + }; + } + + /** + * Format a number of milliseconds to a string of seconds + * + * @param ms - The number of milliseconds + * @returns The number of seconds + */ + static formatMs(ms: number): string { + return `${(ms / 1000).toLocaleString(undefined, { + maximumFractionDigits: 2, + })}s`; + } + + /** + * Update the rate limit information and last response + * + * @param rateLimit - The rate limit information + * @param lastResponse - The last response + */ + protected updateRateLimitInfo( + rateLimit: RateLimitInfo, + lastResponse?: Response, + ): void { + this.limit = rateLimit.limit; + this.remaining = rateLimit.remaining; + this.reset = rateLimit.reset; + this.lastResponse = lastResponse; + } + + /** + * Wrap a `got` call with a rate limit + * + * @param callback - The call to `got`. It should return a Response, so do NOT chain `.json()` or `.text()`—use responseType instead. + * @param options - Options for rate limit handling + * @returns The `got` Response object + * @example + * ```ts + * for (const page of [1, 2, 3]) { + * const response = await rateLimiter.withRateLimit(() => + * got.get<{ id: number; title: string }>(`https://example.com/posts?page=${page}`, { responseType: 'json' }), + * ); + * console.log(response.body.title); + * } + * ``` + */ + async withRateLimit( + callback: () => CancelableRequest>, + { + maxWaitTimeMs = RateLimitClient.MAX_WAIT_TIME_MS, + }: { + /** + * The number of milliseconds to wait until the rate limit resets. + * If not provided, the default maximum wait time will be used. + */ + maxWaitTimeMs?: number; + } = {}, + ): Promise> { + if ( + this.reset !== undefined && + this.reset > new Date() && + (this.lastResponse?.statusCode === 429 || + (this.remaining !== undefined && this.remaining <= 0)) + ) { + const timeUntilResetMs = this.reset.getTime() - new Date().getTime(); + + // Throw if it's beyond the maximum wait time + if (timeUntilResetMs > maxWaitTimeMs) { + throw new Error( + `The time until the rate limit resets (${RateLimitClient.formatMs( + timeUntilResetMs, + )})` + + ` is beyond the maximum wait time (${RateLimitClient.formatMs( + maxWaitTimeMs, + )}). Try again in ${this.reset.toLocaleString()}${ + this.limit + ? `, when the rate limit resets to ${this.limit.toLocaleString()}.` + : '' + }.`, + { cause: this.lastResponse?.body }, + ); + } + + // Wait until it resets before trying the request + logger.warn( + `Waiting until the rate limit resets (${RateLimitClient.formatMs( + timeUntilResetMs, + )}).`, + ); + await new Promise((resolve) => { + setTimeout(resolve, timeUntilResetMs); + }); + } + + let response: Response; + try { + response = await callback(); + } catch (error) { + if (error instanceof RequestError && error.response?.statusCode === 429) { + const rateLimit = RateLimitClient.getRateLimitInfo( + error.response.headers, + ); + this.updateRateLimitInfo(rateLimit, error.response); + + logger.warn('A rate limit was exceeded on this request. Retrying...'); + + // If the error is a 429, we can retry the request + return this.withRateLimit(callback); + } + throw error; + } + + const rateLimit = RateLimitClient.getRateLimitInfo(response.headers); + this.updateRateLimitInfo(rateLimit, response); + + return response; + } +} diff --git a/src/lib/tests/rateLimit.test.ts b/src/lib/tests/rateLimit.test.ts new file mode 100644 index 00000000..340f928f --- /dev/null +++ b/src/lib/tests/rateLimit.test.ts @@ -0,0 +1,92 @@ +import { describe, it, beforeEach, expect } from 'vitest'; +import got from 'got'; +import { RateLimitClient } from '../rateLimit'; + +/** + * A mock of the RateLimitClient class for testing + */ +class RateLimitClientMock extends RateLimitClient { + /** + * Update the rate limit information and last response. + * Forcing this to be public to allow testing. + * + * @param args - The arguments to pass to the parent class + */ + public override updateRateLimitInfo( + ...args: Parameters + ): ReturnType { + super.updateRateLimitInfo(...args); + } +} + +describe('rateLimit', () => { + let rateLimiter: RateLimitClientMock; + beforeEach(() => { + rateLimiter = new RateLimitClientMock(); + }); + + it('should wait until the rate limit resets', async () => { + const timeToWait = 1000; + rateLimiter.updateRateLimitInfo({ + limit: 10, + remaining: 0, + reset: new Date(Date.now() + timeToWait), + }); + const before = Date.now(); + try { + await rateLimiter.withRateLimit(() => got.get('does/not/matter')); + } catch { + // Doesn't matter + } + const after = Date.now(); + + expect(after - before).toBeGreaterThanOrEqual(timeToWait); + }); + + it('should return the response from the API', async () => { + const response = await rateLimiter.withRateLimit(() => + got.get<{ + /** The pokemon's height in decimeters */ + height: number; + }>('https://pokeapi.co/api/v2/pokemon/ditto', { + responseType: 'json', + }), + ); + + expect(response.statusCode).toBe(200); + expect(response.body.height).toBe(3); + }); + + it('should decrement rate limit remaining', async () => { + // This public API supports the 'x-ratelimit-remaining' header + const url = 'https://www.reddit.com/r/javascript.json'; + + // Run an initial request to hydrate the rate limit info + const res = await rateLimiter.withRateLimit(() => got.get(url)); + let prevRemaining: number = rateLimiter.remaining!; + expect(rateLimiter.remaining).toBeDefined(); + + // Test corner case: wait for a reset if we're 3 seconds away from resetting the `remaining` counter + const reset = res.headers['x-ratelimit-reset']; + if (typeof reset === 'string') { + // `x-ratelimit-reset` is in seconds on the Reddit API (and not a timestamp like our RateLimitClient expects) + const resetSeconds = Number.parseInt(reset, 10); + if (resetSeconds <= 3) { + // Wait for the reset to happen + await new Promise((resolve) => { + setTimeout(resolve, (resetSeconds + 1) * 1000); + }); + // Re-run the initial request to hydrate the rate limit info + await rateLimiter.withRateLimit(() => got.get(url)); + prevRemaining = rateLimiter.remaining!; + } + } + + // The actual test: confirm that it decrements properly with each request + for (let i = 0; i < 3; i += 1) { + await rateLimiter.withRateLimit(() => got.get(url)); + expect(prevRemaining - 1).toBe(rateLimiter.remaining!); + prevRemaining = rateLimiter.remaining!; + } + }); +});