Optimized sliding-window

src/lua-scripts/hash.ts

@@ -36,11 +36,11 @@ export const SCRIPTS: {
     slidingWindow: {
       limit: {
         script: Single.slidingWindowLimitScript,
-        hash: "e1391e429b699c780eb0480350cd5b7280fd9213"
+        hash: "20408d957f9eabcc56b113f37785bbbbd301329c"
       },
       getRemaining: {
         script: Single.slidingWindowRemainingTokensScript,
-        hash: "65a73ac5a05bf9712903bc304b77268980c1c417"
+        hash: "5584aba50ee1821800228a532bec593cd08b7022"
       },
     },
     tokenBucket: {

src/lua-scripts/single.ts

@@ -25,59 +25,81 @@ export const fixedWindowRemainingTokensScript = `
     `;
 
 export const slidingWindowLimitScript = `
-  local currentKey  = KEYS[1]           -- identifier including prefixes
-  local previousKey = KEYS[2]           -- key of the previous bucket
+  local key         = KEYS[1]           -- identifier including prefixes
   local tokens      = tonumber(ARGV[1]) -- tokens per window
   local now         = ARGV[2]           -- current timestamp in milliseconds
   local window      = ARGV[3]           -- interval in milliseconds
   local incrementBy = ARGV[4]           -- increment rate per request at a given value, default is 1
-
-  local requestsInCurrentWindow = redis.call("GET", currentKey)
-  if requestsInCurrentWindow == false then
-    requestsInCurrentWindow = 0
-  end
-
-  local requestsInPreviousWindow = redis.call("GET", previousKey)
-  if requestsInPreviousWindow == false then
-    requestsInPreviousWindow = 0
+
+  -- following block is equivalent to, but less commands than "GET" + "INCRBY"
+  local value = redis.pcall("INCRBY", key, incrementBy)
+  -- by design, blocked keys hold non-integer string, see below. So above command might throw error
+  if type(value) ~= "number" then return -1 end
+  value = value - incrementBy    -- pre-existing value or 0
+
+  -- extract info by decoding value. To understand the encoding, see newValue definition below
+  local requestsInCurrentWindow = value % tokens
+  local valueByTokens = math.floor(value / tokens)
+  local requestsInPreviousWindow = valueByTokens % tokens
+  local bitRepresentingWindow = math.floor(valueByTokens / tokens)
+  local currWinIndex = math.floor(now / window)
+  local currWinMod2 = currWinIndex % 2
+  local needsReset = bitRepresentingWindow ~= currWinMod2 or value == 0
+
+  local newValue
+  if needsReset then
+    requestsInPreviousWindow = requestsInCurrentWindow
+    requestsInCurrentWindow = incrementBy
+    newValue = (currWinMod2*tokens + requestsInPreviousWindow)*tokens + requestsInCurrentWindow
+  else
+    requestsInCurrentWindow = requestsInCurrentWindow + incrementBy
   end
+
   local percentageInCurrent = ( now % window ) / window
   -- weighted requests to consider from the previous window
   requestsInPreviousWindow = math.floor(( 1 - percentageInCurrent ) * requestsInPreviousWindow)
-  if requestsInPreviousWindow + requestsInCurrentWindow >= tokens then
-    return -1
+  local remaining = tokens - ( requestsInPreviousWindow + requestsInCurrentWindow )
+
+  if remaining <= 0 then
+    local reset = ( currWinIndex + 1 ) * window    -- expire by next window
+    -- set a string at key so that next "INCRBY" throws error
+    redis.call("SET", key, "BLOCKED!", "PXAT", reset)
+  elseif needsReset then
+    local reset = ( currWinIndex + 2 ) * window    -- live for another window
+    redis.call("SET", key, newValue, "PXAT", reset)   -- store value encoding new info
   end
 
-  local newValue = redis.call("INCRBY", currentKey, incrementBy)
-  if newValue == tonumber(incrementBy) then
-    -- The first time this key is set, the value will be equal to incrementBy.
-    -- So we only need the expire command once
-    redis.call("PEXPIRE", currentKey, window * 2 + 1000) -- Enough time to overlap with a new window + 1 second
-  end
-  return tokens - ( newValue + requestsInPreviousWindow )
+  return remaining
 `;
 
 export const slidingWindowRemainingTokensScript = `
-  local currentKey  = KEYS[1]           -- identifier including prefixes
-  local previousKey = KEYS[2]           -- key of the previous bucket
-  local now         = ARGV[1]           -- current timestamp in milliseconds
-  local window      = ARGV[2]           -- interval in milliseconds
-
-  local requestsInCurrentWindow = redis.call("GET", currentKey)
-  if requestsInCurrentWindow == false then
+  local key         = KEYS[1]           -- identifier including prefixes
+  local tokens      = tonumber(ARGV[1]) -- tokens per window
+  local now         = ARGV[2]           -- current timestamp in milliseconds
+  local window      = ARGV[3]           -- interval in milliseconds
+
+  local value = redis.call("GET", key) or 0 -- value is 0 when key doesn't exist
+  -- by design, blocked keys hold non-integer string
+  if tonumber(value) == nil then return 0 end
+
+  -- extract info by decoding value
+  local requestsInCurrentWindow = value % tokens
+  local valueByTokens = math.floor(value / tokens)
+  local requestsInPreviousWindow = valueByTokens % tokens
+  local bitRepresentingWindow = math.floor(valueByTokens / tokens)
+  local currWinIndex = math.floor(now / window)
+  local currWinMod2 = currWinIndex % 2
+  local needsReset = bitRepresentingWindow ~= currWinMod2 or value == 0
+
+  if needsReset then
+    requestsInPreviousWindow = requestsInCurrentWindow
     requestsInCurrentWindow = 0
   end
-
-  local requestsInPreviousWindow = redis.call("GET", previousKey)
-  if requestsInPreviousWindow == false then
-    requestsInPreviousWindow = 0
-  end
-
+
   local percentageInCurrent = ( now % window ) / window
   -- weighted requests to consider from the previous window
   requestsInPreviousWindow = math.floor(( 1 - percentageInCurrent ) * requestsInPreviousWindow)
-
-  return requestsInPreviousWindow + requestsInCurrentWindow
+  return tokens - ( requestsInPreviousWindow + requestsInCurrentWindow )
 `;
 
 export const tokenBucketLimitScript = `

src/single.ts

@@ -260,11 +260,7 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
     return () => ({
       async limit(ctx: RegionContext, identifier: string, rate?: number) {
         const now = Date.now();
-
         const currentWindow = Math.floor(now / windowSize);
-        const currentKey = [identifier, currentWindow].join(":");
-        const previousWindow = currentWindow - 1;
-        const previousKey = [identifier, previousWindow].join(":");
 
         if (ctx.cache) {
           const { blocked, reset } = ctx.cache.isBlocked(identifier);
@@ -285,7 +281,7 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
         const remainingTokens = await safeEval(
           ctx,
           SCRIPTS.singleRegion.slidingWindow.limit,
-          [currentKey, previousKey],
+          [identifier],
           [tokens, now, windowSize, incrementBy],
         ) as number;
 
@@ -306,34 +302,25 @@ export class RegionRatelimit extends Ratelimit<RegionContext> {
       async getRemaining(ctx: RegionContext, identifier: string) {
         const now = Date.now();
         const currentWindow = Math.floor(now / windowSize);
-        const currentKey = [identifier, currentWindow].join(":");
-        const previousWindow = currentWindow - 1;
-        const previousKey = [identifier, previousWindow].join(":");
 
-        const usedTokens = await safeEval(
+        const remaining = await safeEval(
           ctx,
           SCRIPTS.singleRegion.slidingWindow.getRemaining,
-          [currentKey, previousKey],
-          [now, windowSize],
+          [identifier],
+          [tokens, now, windowSize],
         ) as number;
 
         return {
-          remaining: Math.max(0, tokens - usedTokens),
+          remaining,
           reset: (currentWindow + 1) * windowSize
         }
       },
       async resetTokens(ctx: RegionContext, identifier: string) {
-        const pattern = [identifier, "*"].join(":");
         if (ctx.cache) {
           ctx.cache.pop(identifier)
         }
 
-        await safeEval(
-          ctx,
-          RESET_SCRIPT,
-          [pattern],
-          [null],
-        ) as number;
+        await ctx.redis.del(identifier);
       },
     });
   }