fix: quota usage (#65)

* chore: refactor

* fix: increment quota count only if erl is not active right now anyway
TODO: Write a test for this

* fix: return correct previously used quota at first activation

* fix: return correct erl_quota_left

* test: ability to exhaust all erl quota

---------

Co-authored-by: Ece Tavasli <[email protected]>

lib/db.js

@@ -302,7 +302,7 @@ class LimitDBRedis extends EventEmitter {
         ms_per_interval: bucketKeyConfig.ms_per_interval,
         size: bucketKeyConfig.size,
         erl_activation_period_seconds: 900,
-        erl_quota_amount: 0,
+        erl_quota: 0,
         erl_quota_interval: 'quota_per_calendar_month',
         erl_configured_for_bucket: false,
         ...bucketKeyConfig.elevated_limits
@@ -318,7 +318,7 @@ class LimitDBRedis extends EventEmitter {
         elevated_limits.ms_per_interval,
         elevated_limits.size,
         elevated_limits.erl_activation_period_seconds,
-        elevated_limits.erl_quota_amount,
+        elevated_limits.erl_quota,
         erl_quota_expiration,
         elevated_limits.erl_configured_for_bucket ? 1 : 0,
         (err, results) => {
@@ -345,7 +345,7 @@ class LimitDBRedis extends EventEmitter {
               triggered: erl_triggered,
               activated: erl_activate_for_bucket,
               quota_remaining: erl_quota_count,
-              quota_allocated: elevated_limits.erl_quota_amount,
+              quota_allocated: elevated_limits.erl_quota,
               erl_activation_period_seconds: elevated_limits.erl_activation_period_seconds,
             },
           };

lib/take_elevated.lua

@@ -6,7 +6,7 @@ local drip_interval = tonumber(ARGV[5])
 local erl_tokens_per_ms = tonumber(ARGV[6])
 local erl_bucket_size = tonumber(ARGV[7])
 local erl_activation_period_seconds = tonumber(ARGV[8])
-local erl_quota_amount = tonumber(ARGV[9])
+local erl_quota = tonumber(ARGV[9])
 local erl_quota_expiration_epoch = tonumber(ARGV[10])
 local erl_configured_for_bucket = tonumber(ARGV[11]) == 1
 
@@ -47,29 +47,28 @@ local function calculateNewBucketContent(current, tokens_per_ms, bucket_size, cu
     end
 end
 
-local function takeERLQuota(erl_quota_key, erl_quota_amount, erl_quota_expiration_epoch)
-    if erl_quota_amount <= 0 then
+local function takeERLQuota(erl_quota_key, erl_quota, erl_quota_expiration_epoch)
+    if erl_quota <= 0 then
         -- no quota available to take
         return 0
     end
 
-    local get_quota_result = redis.call('GET', erl_quota_key)
-    if type(get_quota_result) ~= 'string' then
+    local previously_used_quota = redis.call('GET', erl_quota_key)
+    if type(previously_used_quota) ~= 'string' then
         -- first activation. Set quota to 1 and return.
         redis.call('SET', erl_quota_key, 1, 'PXAT', string.format('%.0f', erl_quota_expiration_epoch))
-        return 1
+        return 0
     end
 
-    local erl_quota_used = tonumber(get_quota_result)
-    if erl_quota_used >= erl_quota_amount then
-        -- quota is exceeded. Return the current quota.
-        return erl_quota_used
+    previously_used_quota = tonumber(previously_used_quota)
+    if previously_used_quota >= erl_quota then
+        -- quota is already exceeded. Return the current total used quota.
+        return previously_used_quota
     end
 
     -- quota is not exceeded. Increment and return.
-    local new_quota = erl_quota_used + 1
-    redis.call('SET', erl_quota_key, new_quota, 'PXAT', string.format('%.0f', erl_quota_expiration_epoch))
-    return new_quota
+    redis.call('SET', erl_quota_key, previously_used_quota+1, 'PXAT', string.format('%.0f', erl_quota_expiration_epoch))
+    return previously_used_quota
 end
 
 -- Enable verbatim replication to ensure redis sends script's source code to all masters
@@ -98,21 +97,21 @@ if enough_tokens then
     end
 else
     -- if tokens are not enough, see if activating erl will help.
-    if erl_configured_for_bucket and is_erl_activated == 0 and erl_quota_amount > 0 then
+    if erl_configured_for_bucket and is_erl_activated == 0 and erl_quota > 0 then
         local used_tokens = bucket_size - bucket_content_after_refill
         local bucket_content_after_erl_activation = erl_bucket_size - used_tokens
         local enough_tokens_after_erl_activation = bucket_content_after_erl_activation >= tokens_to_take
         if enough_tokens_after_erl_activation then
-            local erl_quota_used = takeERLQuota(erl_quota_key, erl_quota_amount, erl_quota_expiration_epoch)
-            if erl_quota_used < erl_quota_amount then
+            local previously_used_quota = takeERLQuota(erl_quota_key, erl_quota, erl_quota_expiration_epoch)
+            if previously_used_quota < erl_quota then
                 enough_tokens = enough_tokens_after_erl_activation -- we are returning this value, thus setting it
                 bucket_content_after_take = math.min(bucket_content_after_erl_activation - tokens_to_take, erl_bucket_size)
                 -- save erl state
                 redis.call('SET', erlKey, '1')
                 redis.call('EXPIRE', erlKey, erl_activation_period_seconds)
                 is_erl_activated = 1
                 erl_triggered = true
-                erl_quota_left = erl_quota_amount - erl_quota_used
+                erl_quota_left = erl_quota - previously_used_quota - 1
             end
         end
     end

lib/utils.js

@@ -64,11 +64,11 @@ function normalizeElevatedTemporals(params) {
     if (!(intervalShortcut in params)) {
       return;
     }
-    type.erl_quota_amount = params[intervalShortcut];
+    type.erl_quota = params[intervalShortcut];
     type.erl_quota_interval = intervalShortcut;
   });
 
-  if (!type.erl_quota_interval || type.erl_quota_amount === undefined) {
+  if (!type.erl_quota_interval || type.erl_quota === undefined) {
     throw new LimitdRedisConfigurationError('a valid quota amount per interval is required for elevated limits', {code: 202});
   }
 

test/db.tests.js

@@ -1245,6 +1245,85 @@ describe('LimitDBRedis', () => {
           });
       });
 
+      it("should exhaust all monthly erl quota before rate limiting", (done) => {
+        const bucketName = 'test-bucket';
+        db.configurateBucket(bucketName, {
+          size: 1,
+          per_minute: 1,
+          elevated_limits: {
+            size: 3,
+            per_minute: 3,
+            erl_activation_period_seconds: 900,
+            quota_per_calendar_month: 1
+          }
+        });
+        const params = {
+          type: bucketName,
+          key: 'some_key ',
+          elevated_limits: { erl_is_active_key: 'some_erl_active_identifier', erl_quota_key: 'erlquotakey' },
+        };
+
+        // check erl not activated yet
+        redisExistsPromise(params.elevated_limits.erl_is_active_key)
+            .then((erlIsActiveExists) => assert.equal(erlIsActiveExists, 0))
+            // check erl_quota_key does not exist
+            .then(() => redisExistsPromise(params.elevated_limits.erl_quota_key)
+            .then((erl_quota_keyExists) => assert.equal(erl_quota_keyExists, 0)))
+            // attempt to take elevated should work for first token
+            .then(() => takeElevatedPromise(params))
+            .then((result) => {
+              assert.isTrue(result.conformant);
+              assert.isFalse(result.elevated_limits.activated);
+              assert.isFalse(result.elevated_limits.triggered);
+              assert.isTrue(result.elevated_limits.erl_configured_for_bucket)
+              assert.equal(result.limit, 1);
+            })
+            .then(() => redisExistsPromise(params.elevated_limits.erl_is_active_key))
+            .then((erl_is_active_keyExists) => assert.equal(erl_is_active_keyExists, 0))
+            .then(() => redisExistsPromise(params.elevated_limits.erl_quota_key)
+            .then((erl_quota_keyExists) => assert.equal(erl_quota_keyExists, 0)))
+            // next takeElevated should activate ERL
+            .then(() => takeElevatedPromise(params))
+            .then((result) => {
+              assert.isTrue(result.conformant);
+              assert.isTrue(result.elevated_limits.activated);
+              assert.isTrue(result.elevated_limits.triggered);
+              assert.equal(result.limit, 3);
+            })
+            .then(() => redisExistsPromise(params.elevated_limits.erl_is_active_key))
+            .then((erl_is_active_keyExists) => assert.equal(erl_is_active_keyExists, 1))
+            // check erlQuota was increased
+            .then(() => redisGetPromise(params.elevated_limits.erl_quota_key))
+            .then((erl_quota_keyValue) => assert.equal(erl_quota_keyValue, 1))
+            // exhaust the bucket
+            .then(() => takeElevatedPromise(params))
+            .then((result) => {
+              assert.isTrue(result.conformant);
+              assert.isTrue(result.elevated_limits.activated);
+              assert.isFalse(result.elevated_limits.triggered);
+              assert.equal(result.limit, 3);
+            })
+            .then(() => redisGetPromise(params.elevated_limits.erl_quota_key))
+            .then((erl_quota_keyValue) => assert.equal(erl_quota_keyValue, 1))
+            // remove erl_is_active_key to stop ERL
+            .then(() => redisDeletePromise(params.elevated_limits.erl_is_active_key))
+            // next takeElevated should not activate ERL
+            .then(() => takeElevatedPromise(params))
+            .then((result) => {
+              assert.isFalse(result.conformant);
+              assert.isFalse(result.elevated_limits.activated);
+              assert.isFalse(result.elevated_limits.triggered);
+              assert.isTrue(result.elevated_limits.erl_configured_for_bucket);
+              assert.equal(result.limit, 1);
+            })
+            .then(() => redisExistsPromise(params.elevated_limits.erl_is_active_key))
+            .then((erl_is_active_keyExists) => assert.equal(erl_is_active_keyExists, 0))
+            // check erlQuota was NOT increased
+            .then(() => redisExistsPromise(params.elevated_limits.erl_quota_key))
+            .then((erl_quota_keyValue) => assert.equal(erl_quota_keyValue, 1))
+            .then(() => done());
+      });
+
       describe('when erl is activated for the tenant with multiple bucket configurations', () => {
         const nonERLTestBucket = 'nonerl-test-bucket';
         const ERLBucketName = 'erl-test-bucket';

test/utils.tests.js

@@ -34,7 +34,7 @@ describe('utils', () => {
 
       expect(elevated_limits).excluding('drip_interval').to.deep.equal({
         size: 300,
-        erl_quota_amount: 192,
+        erl_quota: 192,
         erl_quota_interval: 'quota_per_calendar_month',
         interval: 1000,
         per_interval: 300,
@@ -101,7 +101,7 @@ describe('utils', () => {
       });
       expect(overrides['127.0.0.1'].elevated_limits).excluding('drip_interval').to.deep.equal({
         erl_activation_period_seconds: 900,
-        erl_quota_amount: 10,
+        erl_quota: 10,
         erl_quota_interval: "quota_per_calendar_month",
         size: 400,
         interval: 1000,