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,12 +55,8 @@ 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);
-  if (requestCount === 1) {
-    // For the first increment, set an expiration to clean up this key.
-    await redis.expire(key, RATE_LIMIT_WINDOW_SECONDS);
-  }
+  // first read the request count from redis
+  const requestCount = Number((await redis.get(key)) || "0");
 
   // Get the limit for this window accounting for the sample rate.
   const limitPerWindow =
@@ -71,6 +73,23 @@ export async function rateLimit(args: {
     };
   }
 
+  // only need to increment the request count if it's not already set
+  // we are setting the request count, however we are not waiting on this to be complete
+  if (requestCount === 0) {
+    // For the first increment, set an expiration to clean up this key (EX).
+    redis
+      .set(key, requestCount + 1, "EX", RATE_LIMIT_WINDOW_SECONDS)
+      ?.catch((err) => {
+        console.error("Failed to increment 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 + 1, "XX")?.catch((err) => {
+      console.error("Failed to increment request count", err);
+    });
+  }
+
   return {
     rateLimited: false,
     requestCount,

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

@@ -3,16 +3,16 @@ import { validServiceConfig, validTeamResponse } from "../../mocks.js";
 import { rateLimit } from "./index.js";
 
 const mockRedis = {
-  incr: vi.fn(),
-  expire: vi.fn(),
+  set: vi.fn(),
+  get: vi.fn(),
 };
 
 describe("rateLimit", () => {
   beforeEach(() => {
     // Clear mock function calls and reset any necessary state.
     vi.clearAllMocks();
-    mockRedis.incr.mockReset();
-    mockRedis.expire.mockReset();
+    mockRedis.set.mockReset();
+    mockRedis.get.mockReset();
   });
 
   afterEach(() => {
@@ -35,7 +35,7 @@ describe("rateLimit", () => {
   });
 
   it("should not rate limit if within limit", async () => {
-    mockRedis.incr.mockResolvedValue(50); // Current count is 50 requests in 10 seconds.
+    mockRedis.get.mockResolvedValue("50"); // Current count is 50 requests in 10 seconds.
 
     const result = await rateLimit({
       team: validTeamResponse,
@@ -50,11 +50,11 @@ describe("rateLimit", () => {
       rateLimit: 50,
     });
 
-    expect(mockRedis.expire).not.toHaveBeenCalled();
+    expect(mockRedis.set).toHaveBeenCalledTimes(1);
   });
 
   it("should rate limit if exceeded hard limit", async () => {
-    mockRedis.incr.mockResolvedValue(51);
+    mockRedis.get.mockResolvedValue(51);
 
     const result = await rateLimit({
       team: validTeamResponse,
@@ -72,11 +72,11 @@ describe("rateLimit", () => {
       errorCode: "RATE_LIMIT_EXCEEDED",
     });
 
-    expect(mockRedis.expire).not.toHaveBeenCalled();
+    expect(mockRedis.set).not.toHaveBeenCalled();
   });
 
   it("expires on the first incr request only", async () => {
-    mockRedis.incr.mockResolvedValue(1);
+    mockRedis.get.mockResolvedValue("1");
 
     const result = await rateLimit({
       team: validTeamResponse,
@@ -90,11 +90,11 @@ describe("rateLimit", () => {
       requestCount: 1,
       rateLimit: 50,
     });
-    expect(mockRedis.expire).toHaveBeenCalled();
+    expect(mockRedis.set).toHaveBeenCalled();
   });
 
   it("enforces rate limit if sampled (hit)", async () => {
-    mockRedis.incr.mockResolvedValue(10);
+    mockRedis.get.mockResolvedValue("10");
     vi.spyOn(global.Math, "random").mockReturnValue(0.08);
 
     const result = await rateLimit({
@@ -117,7 +117,7 @@ describe("rateLimit", () => {
   });
 
   it("does not enforce rate limit if sampled (miss)", async () => {
-    mockRedis.incr.mockResolvedValue(10);
+    mockRedis.get.mockResolvedValue(10);
     vi.spyOn(global.Math, "random").mockReturnValue(0.15);
 
     const result = await rateLimit({