Refactor rate limiting to use separate read/write Redis operations

Author: jnsdls

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

@@ -5,8 +5,13 @@ 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: {
@@ -20,6 +25,7 @@ export async function rateLimit(args: {
    * @default 1.0
    */
   sampleRate?: number;
+  logger?: typeof console;
 }): Promise<RateLimitResult> {
   const { team, limitPerSecond, serviceConfig, redis, sampleRate = 1.0 } = args;
 
@@ -49,11 +55,25 @@ 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 redis.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.