From bd5e13350b614d626d98a61defefbf52301008eb Mon Sep 17 00:00:00 2001 From: Pablo Ubal Naveira Date: Tue, 8 Oct 2024 15:44:57 +0200 Subject: [PATCH] fixed returning modified current_timestamp and reset_ms correctly calculated for fixed window in takeeleavated --- lib/take.lua | 3 ++- lib/take_elevated.lua | 7 +++++-- test/db.tests.js | 7 +++---- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/lib/take.lua b/lib/take.lua index 29b8ab4..0099d65 100644 --- a/lib/take.lua +++ b/lib/take.lua @@ -8,6 +8,7 @@ local fixed_window = tonumber(ARGV[6]) local current_time = redis.call('TIME') local current_timestamp_ms = current_time[1] * 1000 + current_time[2] / 1000 +local redis_timestamp_ms = current_timestamp_ms local current = redis.pcall('HMGET', KEYS[1], 'd', 'r') @@ -55,4 +56,4 @@ elseif drip_interval > 0 then reset_ms = math.ceil(current_timestamp_ms + (bucket_size - new_content) * drip_interval) end -return { new_content, enough_tokens, current_timestamp_ms, reset_ms } +return { new_content, enough_tokens, redis_timestamp_ms, reset_ms } diff --git a/lib/take_elevated.lua b/lib/take_elevated.lua index c3b4810..0337afb 100644 --- a/lib/take_elevated.lua +++ b/lib/take_elevated.lua @@ -30,6 +30,7 @@ end -- get current time from redis, to be used in new bucket size calculations later local current_time = redis.call('TIME') local current_timestamp_ms = current_time[1] * 1000 + current_time[2] / 1000 +local redis_timestamp_ms = current_timestamp_ms local function adjustCurrentTimestampForFixedWindow(current_timestamp_ms) if current[1] and tokens_per_ms then @@ -151,7 +152,9 @@ redis.call('HMSET', lastBucketStateKey, redis.call('EXPIRE', lastBucketStateKey, ttl) local reset_ms = 0 -if drip_interval > 0 then +if fixed_window > 0 then + reset_ms = current_timestamp_ms + fixed_window +elseif drip_interval > 0 then if is_erl_activated == 1 then reset_ms = math.ceil(current_timestamp_ms + (erl_bucket_size - bucket_content_after_take) * drip_interval) else @@ -160,4 +163,4 @@ if drip_interval > 0 then end -- Return the current quota -return { bucket_content_after_take, enough_tokens, current_timestamp_ms, reset_ms, erl_triggered, is_erl_activated, erl_quota_left } +return { bucket_content_after_take, enough_tokens, redis_timestamp_ms, reset_ms, erl_triggered, is_erl_activated, erl_quota_left } \ No newline at end of file diff --git a/test/db.tests.js b/test/db.tests.js index 5408e1a..6172c5e 100644 --- a/test/db.tests.js +++ b/test/db.tests.js @@ -878,14 +878,13 @@ module.exports.tests = (clientCreator) => { describe(`when calling the lua script`, () => { it(`should use fixed window when asked`, (done) => { - const now = Date.now(); const interval = 1000; const key = '123.123.123.123'; takeFunc({ type: 'ip', key, count: 1000, fixed_window: true }, (err, response) => { if (err) return done(err); assert.ok(response.conformant); assert.equal(response.remaining, 0); - assert.closeTo(response.reset, Math.ceil((now + interval) / 1000), 1); + assert.closeTo(response.delta_reset_ms, interval, 100); assert.equal(response.limit, 1000); redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => { const lastDrip = value[0]; @@ -893,7 +892,7 @@ module.exports.tests = (clientCreator) => { takeFunc({ type: 'ip', key, count: 1, fixed_window: true }, (err, response) => { assert.notOk(response.conformant); assert.equal(response.remaining, 0); - assert.closeTo(response.reset, Math.ceil((now + interval) / 1000), 1); + assert.closeTo(response.delta_reset_ms, interval/2, 100); assert.equal(response.limit, 1000); redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => { assert.equal(value[0], lastDrip, 'last drip should not have changed'); @@ -901,7 +900,7 @@ module.exports.tests = (clientCreator) => { takeFunc({ type: 'ip', key, count: 1, fixed_window: true }, (err, response) => { assert.ok(response.conformant); assert.equal(response.remaining, 999); - assert.closeTo(response.reset, Math.ceil((Date.now() + interval) / 1000), 1); + assert.closeTo(response.delta_reset_ms, interval, 100); assert.equal(response.limit, 1000); redisHMGetPromise(`ip:${key}`, ['d', 'r']).then((value) => { assert.notEqual(value[0], lastDrip, 'last drip should have changed');