RateLimitClient

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<unknown> | 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<unknown>,
+  ): 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<TBody = unknown>(
+    callback: () => CancelableRequest<Response<TBody>>,
+    {
+      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<Response<TBody>> {
+    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<TBody>;
+    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;
+  }
+}

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<RateLimitClient['updateRateLimitInfo']>
+  ): ReturnType<RateLimitClient['updateRateLimitInfo']> {
+    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!;
+    }
+  });
+});