add reason field to limit response

can take timeout or cacheBlocked values depending on how the response was returned.
upstash · May 29, 2024 · a695d0c · a695d0c
1 parent 0985f1c
commit a695d0c
Show file tree

Hide file tree

Showing 6 changed files with 28 additions and 13 deletions.
diff --git a/src/cache.test.ts b/src/cache.test.ts
@@ -27,15 +27,19 @@ test("ephemeral cache", async () => {
 
   let passes = 0;
 
+  const reasons: (string | undefined)[] = []
   for (let i = 0; i <= 20; i++) {
-    const { success } = await ratelimit.limit("id");
+    const { success, reason } = await ratelimit.limit("id");
     if (success) {
       passes++;
+    } else {
+      reasons.push(reason)
     }
   }
 
   expect(passes).toBeLessThanOrEqual(10);
   expect(metrics.evalsha).toBe(11);
+  expect(reasons).toContain("cacheBlock")
 
   await new Promise((r) => setTimeout(r, 5000));
 }, 10000);
diff --git a/src/multi.ts b/src/multi.ts
@@ -164,6 +164,7 @@ export class MultiRegionRatelimit extends Ratelimit<MultiRegionContext> {
               remaining: 0,
               reset: reset,
               pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
@@ -354,18 +355,19 @@ export class MultiRegionRatelimit extends Ratelimit<MultiRegionContext> {
 
     return () => ({
       async limit(ctx: MultiRegionContext, identifier: string, rate?: number) {
-        // if (ctx.cache) {
-        //   const { blocked, reset } = ctx.cache.isBlocked(identifier);
-        //   if (blocked) {
-        //     return {
-        //       success: false,
-        //       limit: tokens,
-        //       remaining: 0,
-        //       reset: reset,
-        //       pending: Promise.resolve(),
-        //     };
-        //   }
-        // }
+        if (ctx.cache) {
+          const { blocked, reset } = ctx.cache.isBlocked(identifier);
+          if (blocked) {
+            return {
+              success: false,
+              limit: tokens,
+              remaining: 0,
+              reset: reset,
+              pending: Promise.resolve(),
+              reason: "cacheBlock"
+            };
+          }
+        }
 
         const requestId = randomId();
         const now = Date.now();

diff --git a/src/ratelimit.test.ts b/src/ratelimit.test.ts
@@ -135,6 +135,7 @@ describe("timeout", () => {
     expect(res.limit).toBe(0);
     expect(res.remaining).toBe(0);
     expect(res.reset).toBe(0);
+    expect(res.reason).toBe("timeout");
     expect(duration).toBeGreaterThanOrEqual(900);
     expect(duration).toBeLessThanOrEqual(1100);
 

diff --git a/src/ratelimit.ts b/src/ratelimit.ts
@@ -164,6 +164,7 @@ export abstract class Ratelimit<TContext extends Context> {
                 remaining: 0,
                 reset: 0,
                 pending: Promise.resolve(),
+                reason: "timeout"
               });
             }, this.timeout);
           }),

diff --git a/src/single.ts b/src/single.ts
@@ -158,6 +158,7 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
               remaining: 0,
               reset: reset,
               pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
@@ -265,6 +266,7 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
               remaining: 0,
               reset: reset,
               pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }
@@ -370,6 +372,7 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
               remaining: 0,
               reset: reset,
               pending: Promise.resolve(),
+              reason: "cacheBlock"
             };
           }
         }

diff --git a/src/types.ts b/src/types.ts
@@ -26,6 +26,8 @@ export type RegionContext = {
 };
 export type MultiRegionContext = { regionContexts: Omit<RegionContext[], "cache">; cache?: EphemeralCache };
 
+export type ResponseType = "timeout" | "cacheBlock"
+
 export type Context = RegionContext | MultiRegionContext;
 export type RatelimitResponse = {
   /**
@@ -64,6 +66,8 @@ export type RatelimitResponse = {
    * ```
    */
   pending: Promise<unknown>;
+
+  reason?: ResponseType;
 };
 
 export type Algorithm<TContext> = () => {