feat: add reset field to getRemaining

upstash · Jun 25, 2024 · f2315b6 · f2315b6
1 parent c264a32
commit f2315b6
Show file tree

Hide file tree

Showing 7 changed files with 60 additions and 19 deletions.
diff --git a/src/getRemainingTokens.test.ts b/src/getRemainingTokens.test.ts
@@ -17,22 +17,33 @@ function run<TContext extends Context>(builder: Ratelimit<TContext>) {
       // Stop at any random request call within the limit
       const stopAt = Math.floor(Math.random() * (limit - 1) + 1);
       for (let i = 1; i <= limit; i++) {
-        await builder.limit(id);
+
+        const [limitResult, remainigResult] = await Promise.all([
+          builder.limit(id),
+          builder.getRemaining(id)
+        ])
+        // console.log(limitResult.remaining, remainigResult.remaining);
+
+        expect(limitResult.remaining).toBe(remainigResult.remaining)
+        expect(limitResult.reset).toBe(remainigResult.reset)
         if (i == stopAt) {
           break
         }
       }
 
-      const remaining = await builder.getRemaining(id);
+      const {remaining} = await builder.getRemaining(id);
       expect(remaining).toBe(limit - stopAt);
-    }, 10000);
+    }, {
+      timeout: 10000,
+      retry: 3
+    });
   });
 }
 
 function newRegion(limiter: Algorithm<RegionContext>): Ratelimit<RegionContext> {
   return new RegionRatelimit({
     prefix: crypto.randomUUID(),
-    redis: Redis.fromEnv(),
+    redis: Redis.fromEnv({enableAutoPipelining: true}),
     limiter,
   });
 }
@@ -52,14 +63,17 @@ function newMultiRegion(limiter: Algorithm<MultiRegionContext>): Ratelimit<Multi
       new Redis({
         url: ensureEnv("EU2_UPSTASH_REDIS_REST_URL"),
         token: ensureEnv("EU2_UPSTASH_REDIS_REST_TOKEN"),
+        enableAutoPipelining: true
       }),
       new Redis({
         url: ensureEnv("APN_UPSTASH_REDIS_REST_URL"),
         token: ensureEnv("APN_UPSTASH_REDIS_REST_TOKEN"),
+        enableAutoPipelining: true
       }),
       new Redis({
         url: ensureEnv("US1_UPSTASH_REDIS_REST_URL"),
         token: ensureEnv("US1_UPSTASH_REDIS_REST_TOKEN"),
+        enableAutoPipelining: true
       }),
     ],
     limiter,

diff --git a/src/lua-scripts/single.ts b/src/lua-scripts/single.ts
@@ -124,13 +124,13 @@ export const tokenBucketRemainingTokensScript = `
   local key         = KEYS[1]
   local maxTokens   = tonumber(ARGV[1])
         
-  local bucket = redis.call("HMGET", key, "tokens")
+  local bucket = redis.call("HMGET", key, "refilledAt", "tokens")
 
   if bucket[1] == false then
-    return maxTokens
+    return {maxTokens, -1}
   end
         
-  return tonumber(bucket[1])
+  return {tonumber(bucket[2]), tonumber(bucket[1])}
 `;
 
 export const cachedFixedWindowLimitScript = `

diff --git a/src/multi.ts b/src/multi.ts
@@ -302,7 +302,10 @@ export class MultiRegionRatelimit extends Ratelimit<MultiRegionContext> {
           return accTokens + parsedToken;
         }, 0);
 
-        return Math.max(0, tokens - usedTokens);
+        return {
+          remaining: Math.max(0, tokens - usedTokens),
+          reset: (bucket + 1) * windowDuration
+        };
       },
       async resetTokens(ctx: MultiRegionContext, identifier: string) {
         const pattern = [identifier, "*"].join(":");
@@ -514,7 +517,10 @@ export class MultiRegionRatelimit extends Ratelimit<MultiRegionContext> {
         }));
 
         const usedTokens = await Promise.any(dbs.map((s) => s.request));
-        return Math.max(0, tokens - usedTokens);
+        return {
+          remaining: Math.max(0, tokens - usedTokens),
+          reset: (currentWindow + 1) * windowSize
+        };
       },
       async resetTokens(ctx: MultiRegionContext, identifier: string) {
         const pattern = [identifier, "*"].join(":");

diff --git a/src/ratelimit.ts b/src/ratelimit.ts
@@ -256,7 +256,10 @@ export abstract class Ratelimit<TContext extends Context> {
     await this.limiter().resetTokens(this.ctx, pattern);
   };
 
-  public getRemaining = async (identifier: string): Promise<number> => {
+  public getRemaining = async (identifier: string): Promise<{
+    remaining: number;
+    reset: number;
+  }> => {
     const pattern = [this.prefix, identifier].join(":");
 
     return await this.limiter().getRemaining(this.ctx, pattern);

diff --git a/src/resetUsedTokens.test.ts b/src/resetUsedTokens.test.ts
@@ -25,7 +25,7 @@ function run<TContext extends Context>(builder: Ratelimit<TContext>) {
 
       // reset tokens
       await builder.resetUsedTokens(id);
-      const remaining = await builder.getRemaining(id);
+      const {remaining} = await builder.getRemaining(id);
       expect(remaining).toBe(limit);
     }, 10000);
   });

diff --git a/src/single.ts b/src/single.ts
@@ -214,7 +214,10 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
           [null],
         ) as number;
 
-        return Math.max(0, tokens - usedTokens);
+        return {
+          remaining: Math.max(0, tokens - usedTokens),
+          reset: (bucket + 1) * windowDuration
+        };
       },
       async resetTokens(ctx: RegionContext, identifier: string) {
         const pattern = [identifier, "*"].join(":");
@@ -322,7 +325,10 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
           [now, windowSize],
         ) as number;
 
-        return Math.max(0, tokens - usedTokens);
+        return {
+          remaining: Math.max(0, tokens - usedTokens),
+          reset: (currentWindow + 1) * windowSize
+        }
       },
       async resetTokens(ctx: RegionContext, identifier: string) {
         const pattern = [identifier, "*"].join(":");
@@ -416,15 +422,18 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
       },
       async getRemaining(ctx: RegionContext, identifier: string) {
 
-        const remainingTokens = await safeEval(
+        const [remainingTokens, refilledAt] = await safeEval(
           ctx,
           tokenBucketRemainingTokensScript,
           "getRemainingHash",
           [identifier],
           [maxTokens],
-        ) as number;
+        ) as [number, number];
 
-        return remainingTokens;
+        return {
+          remaining: remainingTokens,
+          reset: refilledAt === -1 ? Date.now() + intervalDuration : refilledAt + intervalDuration
+        };
       },
       async resetTokens(ctx: RegionContext, identifier: string) {
         const pattern = identifier;
@@ -541,7 +550,10 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
         const hit = typeof ctx.cache.get(key) === "number";
         if (hit) {
           const cachedUsedTokens = ctx.cache.get(key) ?? 0;
-          return Math.max(0, tokens - cachedUsedTokens);
+          return {
+            remaining: Math.max(0, tokens - cachedUsedTokens),
+            reset: (bucket + 1) * windowDuration
+          };
         }
 
         const usedTokens = await safeEval(
@@ -551,7 +563,10 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
           [key],
           [null],
         ) as number;
-        return Math.max(0, tokens - usedTokens);
+        return {
+          remaining: Math.max(0, tokens - usedTokens),
+          reset: (bucket + 1) * windowDuration
+        };
       },
       async resetTokens(ctx: RegionContext, identifier: string) {
         // Empty the cache

diff --git a/src/types.ts b/src/types.ts
@@ -100,7 +100,10 @@ export type Algorithm<TContext> = () => {
       cache?: EphemeralCache;
     },
   ) => Promise<RatelimitResponse>;
-  getRemaining: (ctx: TContext, identifier: string) => Promise<number>;
+  getRemaining: (ctx: TContext, identifier: string) => Promise<{
+    remaining: number,
+    reset: number
+  }>;
   resetTokens: (ctx: TContext, identifier: string) => Promise<void>;
 };