From 75f0506fbe7f37489ce2dbd35d777b9c06abdd89 Mon Sep 17 00:00:00 2001 From: Ece Tavasli <138665441+ecita@users.noreply.github.com> Date: Fri, 10 May 2024 19:19:50 +0200 Subject: [PATCH] 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 --- lib/db.js | 6 ++-- lib/take_elevated.lua | 33 +++++++++--------- lib/utils.js | 4 +-- test/db.tests.js | 79 +++++++++++++++++++++++++++++++++++++++++++ test/utils.tests.js | 4 +-- 5 files changed, 102 insertions(+), 24 deletions(-) diff --git a/lib/db.js b/lib/db.js index 19bdcb0..1230fba 100644 --- a/lib/db.js +++ b/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, }, }; diff --git a/lib/take_elevated.lua b/lib/take_elevated.lua index 3aabead..51104fb 100644 --- a/lib/take_elevated.lua +++ b/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,13 +97,13 @@ 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 @@ -112,7 +111,7 @@ else 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 diff --git a/lib/utils.js b/lib/utils.js index f2ad9a2..38002b3 100644 --- a/lib/utils.js +++ b/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}); } diff --git a/test/db.tests.js b/test/db.tests.js index 7eeb6ed..33805cd 100644 --- a/test/db.tests.js +++ b/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'; diff --git a/test/utils.tests.js b/test/utils.tests.js index 843b46f..5868f87 100644 --- a/test/utils.tests.js +++ b/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,