chore: sliding window for rate limit

Author: arcoraven

packages/service-utils/src/core/rateLimit/index.ts

@@ -1,36 +1,28 @@
 import type { CoreServiceConfig, TeamResponse } from "../api.js";
 import type { RateLimitResult } from "./types.js";
 
-const RATE_LIMIT_WINDOW_SECONDS = 10;
+const SLIDING_WINDOW_SECONDS = 10;
 
 // Redis interface compatible with ioredis (Node) and upstash (Cloudflare Workers).
 type IRedis = {
   incr: (key: string) => Promise<number>;
   expire: (key: string, ttlSeconds: number) => Promise<0 | 1>;
+  mget: (keys: string[]) => Promise<(string | null)[]>;
 };
 
+/**
+ * Increments the request count for this team and returns whether the team has hit their rate limit.
+ * Uses a sliding 10 second window.
+ * @param args
+ * @returns
+ */
 export async function rateLimit(args: {
   team: TeamResponse;
   limitPerSecond: number;
   serviceConfig: CoreServiceConfig;
   redis: IRedis;
-  /**
-   * Sample requests to reduce load on Redis.
-   * This scales down the request count and the rate limit threshold.
-   * @default 1.0
-   */
-  sampleRate?: number;
 }): Promise<RateLimitResult> {
-  const { team, limitPerSecond, serviceConfig, redis, sampleRate = 1.0 } = args;
-
-  const shouldSampleRequest = Math.random() < sampleRate;
-  if (!shouldSampleRequest) {
-    return {
-      rateLimited: false,
-      requestCount: 0,
-      rateLimit: 0,
-    };
-  }
+  const { team, limitPerSecond, serviceConfig, redis } = args;
 
   if (limitPerSecond === 0) {
     // No rate limit is provided. Assume the request is not rate limited.
@@ -41,39 +33,54 @@ export async function rateLimit(args: {
     };
   }
 
+  const currentSecond = Math.floor(Date.now() / 1000);
   const serviceScope = serviceConfig.serviceScope;
 
-  // Gets the 10-second window for the current timestamp.
-  const timestampWindow =
-    Math.floor(Date.now() / (1000 * RATE_LIMIT_WINDOW_SECONDS)) *
-    RATE_LIMIT_WINDOW_SECONDS;
-  const key = `rate-limit:${serviceScope}:${team.id}:${timestampWindow}`;
+  // Increment the request count for the current second.
+  const currentKey = getRequestCountAtSecondCacheKey(
+    serviceScope,
+    team.id,
+    currentSecond,
+  );
+  await redis.incr(currentKey);
+  // Expire this key after it's past the sliding window.
+  await redis.expire(currentKey, SLIDING_WINDOW_SECONDS + 1);
 
-  // Increment and get the current request count in this window.
-  const requestCount = await redis.incr(key);
-  if (requestCount === 1) {
-    // For the first increment, set an expiration to clean up this key.
-    await redis.expire(key, RATE_LIMIT_WINDOW_SECONDS);
-  }
+  // Sum the request counts for the last `SLIDING_WINDOW_SECONDS` seconds.
+  const keys = Array.from(
+    { length: SLIDING_WINDOW_SECONDS },
+    (_, i) => `rate-limit:${serviceScope}:${team.id}:${currentSecond - i}`,
+  );
+  const counts = await redis.mget(keys);
+  const totalRequests = counts.reduce(
+    (sum, count) => sum + (count ? Number.parseInt(count) : 0),
+    0,
+  );
 
-  // Get the limit for this window accounting for the sample rate.
-  const limitPerWindow =
-    limitPerSecond * sampleRate * RATE_LIMIT_WINDOW_SECONDS;
+  const limit = limitPerSecond * SLIDING_WINDOW_SECONDS;
 
-  if (requestCount > limitPerWindow) {
+  if (totalRequests > limit) {
     return {
       rateLimited: true,
-      requestCount,
-      rateLimit: limitPerWindow,
+      requestCount: totalRequests,
+      rateLimit: limit,
       status: 429,
-      errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} reqs/sec. To get higher rate limits, contact us at https://thirdweb.com/contact-us.`,
+      errorMessage: `You've exceeded your ${serviceScope} rate limit at ${limitPerSecond} reqs/sec. Please upgrade your plan to get higher rate limits.`,
       errorCode: "RATE_LIMIT_EXCEEDED",
     };
   }
 
   return {
     rateLimited: false,
-    requestCount,
-    rateLimit: limitPerWindow,
+    requestCount: totalRequests,
+    rateLimit: limit,
   };
 }
+
+function getRequestCountAtSecondCacheKey(
+  serviceScope: CoreServiceConfig["serviceScope"],
+  teamId: string,
+  second: number,
+) {
+  return `rate-limit:${serviceScope}:${teamId}:${second}`;
+}