perf: optimize lua

Author: Kikobeats

src/index.js

@@ -21,45 +21,65 @@ module.exports = class Limiter {
         local max = tonumber(ARGV[3])
         local start = now - duration * 1000
 
-        -- Remove expired entries
-        redis.call('zremrangebyscore', key, 0, start)
+        -- Check if the key exists
+        local exists = redis.call('EXISTS', key)
 
-        -- Get current count
-        local count = redis.call('zcard', key)
+        local count = 0
+        local oldest = now
 
-        -- Calculate remaining
-        local remaining = max - count
+        if exists == 1 then
+          -- Remove expired entries based on the current duration
+          redis.call('ZREMRANGEBYSCORE', key, 0, start)
 
-        -- Add current request
-        redis.call('zadd', key, now, now)
+          -- Get count
+          count = redis.call('ZCARD', key)
 
-        -- Optimize: Only fetch oldest entry if we need it
-        local oldest
-        local oldest_result = redis.call('zrange', key, 0, 0)
-        oldest = #oldest_result > 0 and tonumber(oldest_result[1]) or now
+          -- Get oldest timestamp if we have entries
+          if count > 0 then
+            local oldest_result = redis.call('ZRANGE', key, 0, 0)
+            oldest = tonumber(oldest_result[1])
+          end
+        end
 
-        -- Optimize: Only fetch oldestInRange if count is at or above max
-        local oldestInRange = now
-        if count >= max then
-          local oldest_in_range_result = redis.call('zrange', key, -max, -max)
-          oldestInRange = #oldest_in_range_result > 0 and tonumber(oldest_in_range_result[1]) or now
+        -- Calculate remaining (before adding current request)
+        local remaining = max - count
+
+        -- Early return if already at limit
+        if remaining <= 0 then
+          local resetMicro = oldest + duration * 1000
+          return {0, math.floor(resetMicro / 1000000), max}
         end
 
-        -- Calculate reset time
-        local resetMicro = (oldestInRange ~= now and oldestInRange or oldest) + duration * 1000
+        -- Add current request with current timestamp
+        redis.call('ZADD', key, now, now)
+
+        -- Calculate reset time and handle trimming if needed
+        local resetMicro
 
-        -- Optimize: Only trim if necessary
+        -- Only perform trim if we're at or over max (based on count before adding)
         if count >= max then
-          redis.call('zremrangebyrank', key, 0, -(max + 1))
+          -- Get the entry at position -max for reset time calculation
+          local oldest_in_range_result = redis.call('ZRANGE', key, -max, -max)
+          local oldestInRange = oldest
+
+          if #oldest_in_range_result > 0 then
+            oldestInRange = tonumber(oldest_in_range_result[1])
+          end
+
+          -- Trim the set
+          redis.call('ZREMRANGEBYRANK', key, 0, -(max + 1))
+
+          -- Calculate reset time based on the entry at position -max
+          resetMicro = oldestInRange + duration * 1000
+        else
+          -- We're under the limit, use the oldest entry for reset time
+          resetMicro = oldest + duration * 1000
         end
 
-        -- Set expiration
-        redis.call('pexpire', key, duration)
-
-        -- Ensure remaining is never negative
-        if remaining < 0 then remaining = 0 end
+        -- Set expiration using the provided duration
+        redis.call('PEXPIRE', key, duration)
 
-        return {remaining, resetMicro / 1000000, max}
+        return {remaining, math.floor(resetMicro / 1000000), max}
       `
     })
   }