feat: add defineCommand

Author: Kikobeats

README.md

@@ -128,6 +128,17 @@ Default: `this.max`
 
 How long keep records of requests in milliseconds. If provided, it overrides the default `duration` value.
 
+## defineCommand
+
+It provides the command definition so you can load it into any [ioredis](https://github.com/redis/ioredis) instance:
+
+```js
+const Redis = require('ioredis')
+const redis = new Redis(uri, {
+  scripts: { ratelimiter: require('async-ratelimiter').defineCommand }
+})
+```
+
 ## Related
 
 - [express-slow-down](https://github.com/nfriedly/express-slow-down) – Slow down repeated requests; use as an alternative (or addition) to express-rate-limit.

src/index.js

@@ -3,6 +3,77 @@
 const assert = require('assert')
 const microtime = require('./microtime')
 
+const defineCommand = {
+  numberOfKeys: 1,
+  lua: `
+    local key = KEYS[1]
+    local now = tonumber(ARGV[1])
+    local duration = tonumber(ARGV[2])
+    local max = tonumber(ARGV[3])
+    local start = now - duration
+
+    -- Check if the key exists
+    local exists = redis.call('EXISTS', key)
+
+    local count = 0
+    local oldest = now
+
+    if exists == 1 then
+      -- Remove expired entries based on the current duration
+      redis.call('ZREMRANGEBYSCORE', key, 0, start)
+
+      -- Get count
+      count = redis.call('ZCARD', key)
+
+      -- 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
+
+    -- Calculate remaining (before adding current request)
+    local remaining = max - count
+
+    -- Early return if already at limit
+    if remaining <= 0 then
+      local resetMicro = oldest + duration
+      return {0, math.floor(resetMicro / 1000), max}
+    end
+
+    -- Add current request with current timestamp
+    redis.call('ZADD', key, now, now)
+
+    -- Calculate reset time and handle trimming if needed
+    local resetMicro
+
+    -- Only perform trim if we're at or over max (based on count before adding)
+    if count >= max then
+      -- 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
+    else
+      -- We're under the limit, use the oldest entry for reset time
+      resetMicro = oldest + duration
+    end
+
+    -- Set expiration using the provided duration
+    redis.call('PEXPIRE', key, duration)
+
+    return {remaining, math.floor(resetMicro / 1000), max}
+  `
+}
+
 module.exports = class Limiter {
   constructor ({ id, db, max = 2500, duration = 3600000, namespace = 'limit' }) {
     assert(db, 'db required')
@@ -11,77 +82,9 @@ module.exports = class Limiter {
     this.max = max
     this.duration = duration
     this.namespace = namespace
-
-    this.db.defineCommand('ratelimiter', {
-      numberOfKeys: 1,
-      lua: `
-        local key = KEYS[1]
-        local now = tonumber(ARGV[1])
-        local duration = tonumber(ARGV[2])
-        local max = tonumber(ARGV[3])
-        local start = now - duration
-
-        -- Check if the key exists
-        local exists = redis.call('EXISTS', key)
-
-        local count = 0
-        local oldest = now
-
-        if exists == 1 then
-          -- Remove expired entries based on the current duration
-          redis.call('ZREMRANGEBYSCORE', key, 0, start)
-
-          -- Get count
-          count = redis.call('ZCARD', key)
-
-          -- 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
-
-        -- Calculate remaining (before adding current request)
-        local remaining = max - count
-
-        -- Early return if already at limit
-        if remaining <= 0 then
-          local resetMicro = oldest + duration
-          return {0, math.floor(resetMicro / 1000), max}
-        end
-
-        -- Add current request with current timestamp
-        redis.call('ZADD', key, now, now)
-
-        -- Calculate reset time and handle trimming if needed
-        local resetMicro
-
-        -- Only perform trim if we're at or over max (based on count before adding)
-        if count >= max then
-          -- 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
-        else
-          -- We're under the limit, use the oldest entry for reset time
-          resetMicro = oldest + duration
-        end
-
-        -- Set expiration using the provided duration
-        redis.call('PEXPIRE', key, duration)
-
-        return {remaining, math.floor(resetMicro / 1000), max}
-      `
-    })
+    if (!this.db.ratelimiter) {
+      this.db.defineCommand('ratelimiter', defineCommand)
+    }
   }
 
   async get ({ id = this.id, max = this.max, duration = this.duration } = {}) {
@@ -103,3 +106,5 @@ module.exports = class Limiter {
     }
   }
 }
+
+module.exports.defineCommand = defineCommand

test/index.js

@@ -9,8 +9,11 @@ const { setTimeout } = require('timers/promises')
 const RateLimiter = require('..')
 
 ;['ioredis'].forEach(function (redisModuleName) {
-  const redisModule = require(redisModuleName)
-  const db = require(redisModuleName).createClient()
+  const RedisModule = require(redisModuleName)
+  const db = new RedisModule('127.0.0.1:6379', {
+    lazyConnect: true,
+    enableAutoPipelining: true
+  })
 
   describe('Limiter with ' + redisModuleName, function () {
     beforeEach(async function () {
@@ -176,7 +179,10 @@ const RateLimiter = require('..')
             duration: 10000,
             max,
             id: 'something',
-            db: redisModule.createClient()
+            db: new RedisModule('127.0.0.1:6379', {
+              lazyConnect: true,
+              enableAutoPipelining: true
+            })
           })
         )
       }
@@ -229,7 +235,10 @@ const RateLimiter = require('..')
         duration: 10000,
         max,
         id: 'asyncsomething',
-        db: redisModule.createClient()
+        db: new RedisModule('127.0.0.1:6379', {
+          lazyConnect: true,
+          enableAutoPipelining: true
+        })
       })
 
       it('should set the count properly without race conditions', async function () {