Skip to content

Commit

Permalink
[refactor] Sending limit_keys as keys, simpler lua script
Browse files Browse the repository at this point in the history
  • Loading branch information
didierofrivia committed Apr 3, 2024
1 parent 5058c14 commit c3a0757
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 19 deletions.
14 changes: 5 additions & 9 deletions limitador/src/storage/redis/redis_async.rs
Original file line number Diff line number Diff line change
Expand Up @@ -286,18 +286,14 @@ impl AsyncRedisStorage {

for (counter, delta) in counters_and_deltas {
script_invocation.key(key_for_counter(&counter));
script_invocation.arg(key_for_counters_of_limit(&counter.limit()));
script_invocation.arg(&counter.seconds());
script_invocation.key(key_for_counters_of_limit(counter.limit()));
script_invocation.arg(counter.seconds());
script_invocation.arg(delta);
}

async {
script_invocation
.invoke_async::<_, _>(&mut con)
.await
}
.instrument(span)
.await?;
async { script_invocation.invoke_async::<_, _>(&mut con).await }
.instrument(span)
.await?;

Ok(())
}
Expand Down
19 changes: 9 additions & 10 deletions limitador/src/storage/redis/scripts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,23 +19,22 @@ pub const SCRIPT_UPDATE_COUNTER: &str = "
end
return c";

// KEYS: List of counter keys
// ARGV[j]: Limit keys
// ARGV[j+1]: TTLs
// ARGV[j+2]: Deltas
// KEY[i]: Counter key
// KEY[i+1]: Limit key
// ARGV[i]: TTLs
// ARGV[i+1]: Deltas
pub const BATCH_UPDATE_COUNTERS: &str = "
local j = 1
for i, counter_key in ipairs(KEYS) do
local limit_key = ARGV[j]
local ttl = tonumber(ARGV[j+1])
local delta = tonumber(ARGV[j+2])
for i = 1, #KEYS, 2 do
local counter_key = KEYS[i]
local limit_key = KEYS[i+1]
local ttl = tonumber(ARGV[i])
local delta = tonumber(ARGV[i+1])
local c = redis.call('incrby', counter_key, delta)
if c == delta then
redis.call('expire', counter_key, ttl)
redis.call('sadd', limit_key, counter_key)
end
j = j + 3
end
";

Expand Down

0 comments on commit c3a0757

Please sign in to comment.