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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
|
-- This script updates a token bucket rate limiter with dynamic rate and burst multipliers in Redis.
-- KEYS: Input parameters
-- KEYS[1] - prefix: The Redis key prefix used to store the bucket information.
-- KEYS[2] - now: The current time in milliseconds.
-- KEYS[3] - dynamic_rate_multiplier: A multiplier to adjust the rate limit dynamically.
-- KEYS[4] - dynamic_burst_multiplier: A multiplier to adjust the burst limit dynamically.
-- KEYS[5] - max_dyn_rate: The maximum allowed value for the dynamic rate multiplier.
-- KEYS[6] - max_burst_rate: The maximum allowed value for the dynamic burst multiplier.
-- KEYS[7] - expire: The expiration time for the Redis key storing the bucket information, in seconds.
-- KEYS[8] - number_of_recipients: The number of requests to be allowed (or the increase rate).
-- KEYS[9] - Enable dynamic ratelimits
-- 1. Retrieve the last hit time and initialize variables
local prefix = KEYS[1]
local last = redis.call('HGET', prefix, 'l')
local now = tonumber(KEYS[2])
local nrcpt = tonumber(KEYS[8])
local enable_dynamic = KEYS[9] == 'true'
if not last then
-- 2. Initialize a new bucket if the last hit time is not found (must not happen)
redis.call('HMSET', prefix, 'l', tostring(now), 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0')
redis.call('EXPIRE', prefix, KEYS[7])
return { 1, 1, 1 }
end
-- 3. Update the dynamic rate multiplier based on input parameters
local dr, db = 1.0, 1.0
if enable_dynamic then
local max_dr = tonumber(KEYS[5])
if max_dr > 1 then
local rate_mult = tonumber(KEYS[3])
dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000
if rate_mult > 1.0 and dr < max_dr then
dr = dr * rate_mult
if dr > 0.0001 then
redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
else
redis.call('HSET', prefix, 'dr', '1')
end
elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then
dr = dr * rate_mult
if dr > 0.0001 then
redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
else
redis.call('HSET', prefix, 'dr', '1')
end
end
end
end
-- 4. Update the dynamic burst multiplier based on input parameters
if enable_dynamic then
local max_db = tonumber(KEYS[6])
if max_db > 1 then
local rate_mult = tonumber(KEYS[4])
db = tonumber(redis.call('HGET', prefix, 'db')) / 10000
if rate_mult > 1.0 and db < max_db then
db = db * rate_mult
if db > 0.0001 then
redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
else
redis.call('HSET', prefix, 'db', '1')
end
elseif rate_mult < 1.0 and db > (1.0 / max_db) then
db = db * rate_mult
if db > 0.0001 then
redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
else
redis.call('HSET', prefix, 'db', '1')
end
end
end
end
-- 5. Update the burst and pending values based on the number of recipients (requests)
local burst, pending = unpack(redis.call('HMGET', prefix, 'b', 'p'))
burst, pending = tonumber(burst or '0'), tonumber(pending or '0')
if burst < 0 then
burst = nrcpt
else
burst = burst + nrcpt
end
if pending < nrcpt then
pending = 0
else
pending = pending - nrcpt
end
-- 6. Set the updated values back to Redis and update the expiration time for the bucket
redis.call('HMSET', prefix, 'b', tostring(burst), 'p', tostring(pending), 'l', KEYS[2])
redis.call('EXPIRE', prefix, KEYS[7])
-- 7. Return the updated burst value, dynamic rate multiplier, and dynamic burst multiplier
return { tostring(burst), tostring(dr), tostring(db) }
|
