Merge pull request #50 from auth0/skip-take-inaccuracy

fix: [PSERV-2110] skip_n_calls take amount accounts for current call
auth0 · Dec 12, 2023 · 8e92fbf · 8e92fbf
2 parents 10cba5c + 6cf966e
commit 8e92fbf
Show file tree

Hide file tree

Showing 2 changed files with 136 additions and 29 deletions.
diff --git a/lib/db.js b/lib/db.js
@@ -235,7 +235,8 @@ class LimitDBRedis extends EventEmitter {
         // if count=3, and we go every 5 times, take should 15
         // This parameter is most likely 1, and doing times is an overkill but better safe than sorry.
         if (shouldGoToRedis) {
-          count *= bucketKeyConfig.skip_n_calls;
+          // we need to account for the skipped calls + the current call
+          count *= (bucketKeyConfig.skip_n_calls + 1);
         }
       }
 

diff --git a/test/db.tests.js b/test/db.tests.js
@@ -65,6 +65,16 @@ const buckets = {
         skip_n_calls: 2,
         size: 3,
         per_hour: 3
+      },
+      skipOneSize10: {
+        skip_n_calls: 1,
+        size: 10,
+        per_hour: 0
+      },
+      skipOneSize3: {
+        skip_n_calls: 1,
+        size: 3,
+        per_hour: 0
       }
     }
 
@@ -538,35 +548,131 @@ describe('LimitDBRedis', () => {
       });
     });
 
-    it('should skip calls', (done) => {
-      const params = { type: 'global',  key: 'skipit'};
-
-      async.series([
-        (cb) => db.take(params, cb), // redis
-        (cb) => db.take(params, cb), // cache
-        (cb) => db.take(params, cb), // cache
-        (cb) => {
-          assert.equal(db.callCounts.get('global:skipit').count, 2);
-          cb();
-        },
-        (cb) => db.take(params, cb), // redis
-        (cb) => db.take(params, cb), // cache
-        (cb) => db.take(params, cb), // cache
-        (cb) => db.take(params, cb), // redis (first nonconformant)
-        (cb) => db.take(params, cb), // cache (first cached)
-        (cb) => {
-          assert.equal(db.callCounts.get('global:skipit').count, 1);
-          assert.notOk(db.callCounts.get('global:skipit').res.conformant);
-          cb();
-        },
-      ], (err, _results) => {
-        if (err) {
-          return done(err);
-        }
+    describe('skip calls', () => {
+      it('should skip calls', (done) => {
+        const params = { type: 'global',  key: 'skipit'};
+
+        async.series([
+          (cb) => db.take(params, cb), // redis
+          (cb) => db.take(params, cb), // cache
+          (cb) => db.take(params, cb), // cache
+          (cb) => {
+            assert.equal(db.callCounts.get('global:skipit').count, 2);
+            cb();
+          },
+          (cb) => db.take(params, cb), // redis
+          (cb) => db.take(params, cb), // cache
+          (cb) => db.take(params, cb), // cache
+          (cb) => db.take(params, cb), // redis (first nonconformant)
+          (cb) => db.take(params, cb), // cache (first cached)
+          (cb) => {
+            assert.equal(db.callCounts.get('global:skipit').count, 1);
+            assert.notOk(db.callCounts.get('global:skipit').res.conformant);
+            cb();
+          },
+        ], (err, _results) => {
+          if (err) {
+            return done(err);
+          }
 
-        done();
-      })
-    });
+          done();
+        })
+      });
+
+      it('should take correct number of tokens for skipped calls with single count', (done) => {
+        const params = { type: 'global',  key: 'skipOneSize3'};
+
+        // size = 3
+        // skip_n_calls = 1
+        // no refill
+        async.series([
+          (cb) => db.get(params, (_, {remaining}) => { assert.equal(remaining, 3); cb(); }),
+
+          // call 1 - redis
+          // takes 1 token
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 2);
+            assert.ok(conformant)
+            cb();
+          }),
+
+          // call 2 - skipped
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 2);
+            assert.ok(conformant)
+            cb();
+          }),
+
+          // call 3 - redis
+          // takes 2 tokens here, 1 for current call and one for previously skipped call
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 0);
+            assert.ok(conformant)
+            cb();
+          }),
+
+          // call 4 - skipped
+          // Note: this is the margin of error introduced by skip_n_calls. Without skip_n_calls, this call would be
+          // non-conformant.
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 0);
+            assert.ok(conformant);
+            cb();
+          }),
+
+          // call 5 - redis
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 0);
+            assert.notOk(conformant);
+            cb();
+          }),
+        ], (err, _results) => {
+          if (err) {
+            return done(err);
+          }
+          done();
+        })
+      });
+
+      it('should take correct number of tokens for skipped calls with multi count', (done) => {
+        const params = { type: 'global',  key: 'skipOneSize10', count: 2};
+
+        // size = 10
+        // skip_n_calls = 1
+        // no refill
+        async.series([
+          (cb) => db.get(params, (_, {remaining}) => { assert.equal(remaining, 10); cb(); }),
+
+          // call 1 - redis
+          // takes 2 tokens
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 8);
+            assert.ok(conformant)
+            cb();
+          }),
+
+          // call 2 - skipped
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 8);
+            assert.ok(conformant)
+            cb();
+          }),
+
+          // call 3 - redis
+          // takes 4 tokens here, 2 for current call and 2 for previously skipped call
+          (cb) => db.take(params, (_, { remaining, conformant }) => {
+            assert.equal(remaining, 4);
+            assert.ok(conformant)
+            cb();
+          }),
+        ], (err, _results) => {
+          if (err) {
+            return done(err);
+          }
+          done();
+        })
+      });
+    })
   });
 
   describe('PUT', () => {