Merge pull request #48 from auth0/POC-go-every-n

feat: relax take count's
auth0 · Oct 24, 2023 · e851446 · e851446
2 parents 66eff99 + 05a43d7
commit e851446
Show file tree

Hide file tree

Showing 4 changed files with 80 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -52,6 +52,7 @@ Buckets:
 - `per_interval` (number): is the amount of tokens that the bucket receive on every interval.
 - `interval` (number): defines the interval in milliseconds.
 - `unlimited` (boolean = false): unlimited requests (skip take).
+- `skip_n_calls` (number): take will go to redis every `n` calls instead of going in every take. 
 
 Ping:
 

diff --git a/lib/db.js b/lib/db.js
@@ -2,6 +2,7 @@ const ms = require('ms');
 const fs = require('fs');
 const _ = require('lodash');
 const async = require('async');
+const LRU = require('lru-cache');
 const utils = require('./utils');
 const Redis = require('ioredis');
 const { validateParams } = require('./validation');
@@ -51,6 +52,7 @@ class LimitDBRedis extends EventEmitter {
     this.configurateBuckets(config.buckets);
     this.prefix = config.prefix;
     this.globalTTL = (config.globalTTL || ms('7d')) / 1000;
+    this.callCounts = new LRU({ max: 50 });
 
     const redisOptions = {
       // a low commandTimeout value would likely cause sharded clusters to fail `enableReadyCheck` due to it running `CLUSTER INFO`
@@ -199,7 +201,7 @@ class LimitDBRedis extends EventEmitter {
 
     const key = `${params.type}:${params.key}`;
 
-    const count = this._determineCount({
+    let count = this._determineCount({
       paramsCount: params.count,
       defaultCount: 1,
       bucketKeyConfigSize: bucketKeyConfig.size,
@@ -215,6 +217,31 @@ class LimitDBRedis extends EventEmitter {
       });
     }
 
+    if (bucketKeyConfig.skip_n_calls > 0) {
+      const prevCall = this.callCounts.get(key);
+
+      if (prevCall) {
+        const shouldGoToRedis = prevCall?.count >= bucketKeyConfig.skip_n_calls
+
+
+        if (!shouldGoToRedis) {
+          prevCall.count ++;
+          return process.nextTick(callback, null, prevCall.res);
+        }
+
+        // if lastCall not exists it's the first time that we go to redis.
+        // so we don't change count; subsequently calls take count should be
+        // proportional to the number of call that we skip.
+        // 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;
+        }
+      }
+
+
+    }
+
     this.redis.take(key,
       bucketKeyConfig.ms_per_interval || 0,
       bucketKeyConfig.size,
@@ -238,6 +265,10 @@ class LimitDBRedis extends EventEmitter {
           delayed: false,
         };
 
+        if (bucketKeyConfig.skip_n_calls > 0) {
+          this.callCounts.set(key, { res, count: 0 });
+        }
+
         return callback(null, res);
       });
   }

diff --git a/lib/utils.js b/lib/utils.js
@@ -17,6 +17,7 @@ function normalizeTemporals(params) {
     'interval',
     'size',
     'unlimited',
+    'skip_n_calls'
   ]);
 
   INTERVAL_SHORTCUTS.forEach(intervalShortcut => {

diff --git a/test/db.tests.js b/test/db.tests.js
@@ -57,21 +57,17 @@ const buckets = {
     size: 1,
     per_second: 1
   },
-  cached: {
+  global: {
     size: 3,
     per_hour: 2,
-    enable_cache: true,
     overrides: {
-      faster: {
+      skipit: {
+        skip_n_calls: 2,
         size: 3,
-        per_second: 1,
-        enable_cache: true
-      },
-      disabled: {
-        size: 5,
-        per_hour: 2,
+        per_hour: 3
       }
     }
+
   }
 };
 
@@ -530,6 +526,47 @@ describe('LimitDBRedis', () => {
       });
     });
 
+    it('should call redis and not set local cache count', (done) => {
+      const params = { type: 'global',  key: 'aTenant'};
+      db.take(params, (err) => {
+        if (err) {
+          return done(err);
+        }
+
+        assert.equal(db.callCounts['global:aTenant'], undefined);
+        done();
+      });
+    });
+
+    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();
+      })
+    });
   });
 
   describe('PUT', () => {