feat(redis): add delayMultiplierByGroupEnabled option

animir · Apr 22, 2024 · c2d4c86 · c2d4c86
1 parent b46d082
commit c2d4c86
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 6 deletions.
diff --git a/lib/RateLimiterRedis.js b/lib/RateLimiterRedis.js
@@ -1,12 +1,19 @@
 const RateLimiterStoreAbstract = require('./RateLimiterStoreAbstract');
 const RateLimiterRes = require('./RateLimiterRes');
 
-const incrTtlLuaScript = `redis.call('set', KEYS[1], 0, 'EX', ARGV[2], 'NX') \
+const incrTtlLuaScript = `local ok = redis.call('set', KEYS[1], 0, 'EX', ARGV[2], 'NX') \
 local consumed = redis.call('incrby', KEYS[1], ARGV[1]) \
 local ttl = redis.call('pttl', KEYS[1]) \
 if ttl == -1 then \
   redis.call('expire', KEYS[1], ARGV[2]) \
   ttl = 1000 * ARGV[2] \
+else \
+  local maxPoints = tonumber(ARGV[3]) \
+  if maxPoints > 0 and (consumed-1) % maxPoints == 0 and not ok then \
+    local expireTime = ttl + tonumber(ARGV[2]) * 1000 \
+    redis.call('pexpire', KEYS[1], expireTime) \
+    return {consumed, expireTime} \
+  end \
 end \
 return {consumed, ttl} \
 `;
@@ -27,6 +34,7 @@ class RateLimiterRedis extends RateLimiterStoreAbstract {
     this.client = opts.storeClient;
 
     this._rejectIfRedisNotReady = !!opts.rejectIfRedisNotReady;
+    this._delayMultiplierByGroupEnabled = !!opts.delayMultiplierByGroupEnabled;
 
     this.useRedisPackage = opts.useRedisPackage || this.client.constructor.name === 'Commander' || false;
     this.useRedis3AndLowerPackage = opts.useRedis3AndLowerPackage;
@@ -105,7 +113,7 @@ class RateLimiterRedis extends RateLimiterStoreAbstract {
     if (secDuration > 0) {
       if(!this.useRedisPackage && !this.useRedis3AndLowerPackage){
         return this.client.rlflxIncr(
-          [rlKey].concat([String(points), String(secDuration)]));
+          [rlKey].concat([String(points), String(secDuration), String(this._delayMultiplierByGroupEnabled?this.points:0)]));
       }
       if (this.useRedis3AndLowerPackage) {
         return new Promise((resolve, reject) => {
@@ -118,15 +126,15 @@ class RateLimiterRedis extends RateLimiterStoreAbstract {
           };
 
           if (typeof this.client.rlflxIncr === 'function') {
-            this.client.rlflxIncr(rlKey, points, secDuration, incrCallback);
+            this.client.rlflxIncr(rlKey, points, secDuration, this._delayMultiplierByGroupEnabled?this.points:0, incrCallback);
           } else {
-            this.client.eval(incrTtlLuaScript, 1, rlKey, points, secDuration, incrCallback);
+            this.client.eval(incrTtlLuaScript, 1, rlKey, points, secDuration, this._delayMultiplierByGroupEnabled?this.points:0, incrCallback);
           }
         });
       } else {
         return this.client.eval(incrTtlLuaScript, {
           keys: [rlKey],
-          arguments: [String(points), String(secDuration)],
+          arguments: [String(points), String(secDuration), String(this._delayMultiplierByGroupEnabled?this.points:0)],
         });
       }
     } else {

diff --git a/lib/index.d.ts b/lib/index.d.ts
@@ -255,6 +255,7 @@ interface IRateLimiterRedisOptions extends IRateLimiterStoreOptions {
     rejectIfRedisNotReady?: boolean;
     useRedisPackage?: boolean;
     useRedis3AndLowerPackage?: boolean;
+    delayMultiplierByGroupEnabled?: boolean;
 }
 
 interface ICallbackReady {

diff --git a/test/RateLimiterRedis.ioredis.test.js b/test/RateLimiterRedis.ioredis.test.js
@@ -6,7 +6,7 @@ const sinon = require('sinon');
 const RateLimiterRedis = require('../lib/RateLimiterRedis');
 const Redis = require("ioredis");
 
-describe('RateLimiterRedis with fixed window', function RateLimiterRedisTest() {
+describe.only('RateLimiterRedis with fixed window', function RateLimiterRedisTest() {
   this.timeout(5500);
   let redisMockClient;
 
@@ -59,6 +59,33 @@ describe('RateLimiterRedis with fixed window', function RateLimiterRedisTest() {
       });
   });
 
+  it.only('rejected when consume more than maximum points and delayMultiplierByGroupEnabled', (done) => {
+    const testKey = 'consume2';
+    const rateLimiter = new RateLimiterRedis({
+      storeClient: redisMockClient,
+      points: 1,
+      duration: 5,
+      delayMultiplierByGroupEnabled: true
+    });
+    rateLimiter
+      .consume(testKey)
+      .then(() => {
+        rateLimiter
+          .consume(testKey)
+          .then((res) => {
+            expect(res.msBeforeNext == 5000).to.equal(true);
+            done();
+          })
+          .catch((rejRes) => {
+            expect(rejRes.msBeforeNext >= 5000).to.equal(true);    
+            done();
+          });
+      })
+      .catch((err) => {
+        done(err);
+      });
+  });
+
   it('execute evenly over duration', (done) => {
     const testKey = 'consumeEvenly';
     const rateLimiter = new RateLimiterRedis({