forked from teambition/ratelimiter-go
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathratelimiter.lua
58 lines (46 loc) · 1.34 KB
/
ratelimiter.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
-- KEYS[1] target hash key
-- ARGV[n >= 3] current timestamp, max count, duration, max count, duration, ...
-- HASH: KEYS[1]
-- field:ct(count)
-- field:lt(limit)
-- field:dn(duration)
-- field:rt(reset)
local res = {}
local policyCount = (#ARGV - 1) / 2
local statusKey = '{' .. KEYS[1] .. '}:S'
local limit = redis.call('hmget', KEYS[1], 'ct', 'lt', 'dn', 'rt')
if limit[1] then
res[1] = tonumber(limit[1]) - 1
res[2] = tonumber(limit[2])
res[3] = tonumber(limit[3]) or ARGV[3]
res[4] = tonumber(limit[4])
if policyCount > 1 and res[1] == -1 then
redis.call('incr', statusKey)
redis.call('pexpire', statusKey, res[3] * 2)
local index = tonumber(redis.call('get', statusKey))
if index == 1 then
redis.call('incr', statusKey)
end
end
if res[1] >= -1 then
redis.call('hincrby', KEYS[1], 'ct', -1)
else
res[1] = -1
end
else
local index = 1
if policyCount > 1 then
index = tonumber(redis.call('get', statusKey)) or 1
if index > policyCount then
index = policyCount
end
end
local total = tonumber(ARGV[index * 2])
res[1] = total - 1
res[2] = total
res[3] = tonumber(ARGV[index * 2 + 1])
res[4] = tonumber(ARGV[1]) + res[3]
redis.call('hmset', KEYS[1], 'ct', res[1], 'lt', res[2], 'dn', res[3], 'rt', res[4])
redis.call('pexpire', KEYS[1], res[3])
end
return res