Refactor rate limiting to use separate read/write Redis operations

Author: jnsdls

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

@@ -5,23 +5,40 @@ const RATE_LIMIT_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>;
+  get: (key: string) => Promise<string | null>;
+  set(
+    key: string,
+    value: string | number,
+    secondsToken?: "EX" | "XX",
+    seconds?: number | string,
+  ): Promise<"OK">;
 };
 
 export async function rateLimit(args: {
   team: TeamResponse;
   limitPerSecond: number;
   serviceConfig: CoreServiceConfig;
   redis: IRedis;
+  readRedis?: IRedis;
   /**
    * Sample requests to reduce load on Redis.
    * This scales down the request count and the rate limit threshold.
    * @default 1.0
    */
   sampleRate?: number;
+  logger?: typeof console;
 }): Promise<RateLimitResult> {
-  const { team, limitPerSecond, serviceConfig, redis, sampleRate = 1.0 } = args;
+  const {
+    team,
+    limitPerSecond,
+    serviceConfig,
+    redis,
+    readRedis,
+    sampleRate = 1.0,
+  } = args;
+
+  // fall back to the write redis if read redis is not provided explicitly
+  const readOnlyRedis = readRedis ?? redis;
 
   const shouldSampleRequest = Math.random() < sampleRate;
   if (!shouldSampleRequest) {
@@ -49,11 +66,27 @@ export async function rateLimit(args: {
     RATE_LIMIT_WINDOW_SECONDS;
   const key = `rate-limit:${serviceScope}:${team.id}:${timestampWindow}`;
 
-  // Increment and get the current request count in this window.
-  const requestCount = await redis.incr(key);
+  // first read the request count from redis
+  const currentRequestCountFromRedis = Number(
+    (await readOnlyRedis.get(key)) || "0",
+  );
+
+  const requestCount = currentRequestCountFromRedis + 1;
+
+  // we are setting the request count, however we are not waiting on this to be complete
   if (requestCount === 1) {
-    // For the first increment, set an expiration to clean up this key.
-    await redis.expire(key, RATE_LIMIT_WINDOW_SECONDS);
+    // For the first increment, set an expiration to clean up this key (EX).
+    redis
+      .set(key, requestCount, "EX", RATE_LIMIT_WINDOW_SECONDS)
+      .catch((err) => {
+        console.error("failed to set request count", err);
+      });
+  } else {
+    // For all other increments, just increment the request count.
+    // only set it if it already exists (XX)
+    redis.set(key, requestCount, "XX").catch((err) => {
+      console.error("failed to increment request count", err);
+    });
   }
 
   // Get the limit for this window accounting for the sample rate.