算法原理

系统以恒定的速率产生令牌,然后把令牌放到令牌桶中,令牌桶有一个容量,当令牌桶满了的时候,再向其中放入的令牌会被丢弃;当想要处理一个请求的时候,需要从令牌桶中取出一个令牌,如果此时令牌桶中没有令牌,则拒绝该请求。

数据结构

采用Hash结构存储,字段定义如下:

名称 含义
capacity 令牌桶容量
remain 剩余令牌数
period 时间窗口大小(秒)
quota 时间窗口内的限额
timestamp 生成令牌的时间戳(秒)

脚本代码

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


redis.replicate_commands()

local key = KEYS[1] -- 令牌桶标识
local capacity = tonumber(ARGV[1]) -- 最大容量
local quota = tonumber(ARGV[2]) -- 时间窗口内的限额
local period = tonumber(ARGV[3]) -- 时间窗口大小(秒)
local quantity = tonumber(ARGV[4]) or 1 -- 需要的令牌数量,默认为1
local timestamp = tonumber(redis.call('time')[1]) -- 当前时间戳

assert(type(capacity) == "number", "capacity is not a number!")
assert(type(quota) == "number", "quota is not a number!")
assert(type(period) == "number", "period is not a number!")
assert(type(quantity) == "number", "quantity is not a number!")

-- 第一次请求时创建令牌桶
if (redis.call('exists', key) == 0) then
redis.call('hmset', key, 'remain', capacity, 'timestamp', timestamp)
else
-- 计算从上次生成到现在这段时间应该生成的令牌数
local remain = tonumber(redis.call('hget', key, 'remain'))
local last_reset = tonumber(redis.call('hget', key, 'timestamp'))
local delta_quota = math.floor(((timestamp - last_reset) / period) * quota)
if (delta_quota > 0) then
remain = remain + delta_quota
if (remain > capacity) then
remain = capacity
end
redis.call('hmset', key, 'remain', remain, 'timestamp', timestamp)
end
end

-- 支持动态调整容量和令牌生成速率
redis.call('hmset', key, 'capacity', capacity, 'quota', quota, 'period', period);

local result = {} -- 返回的结果集
local remain = tonumber(redis.call('hget', key, 'remain'))
if (remain < quantity) then
result = {1, capacity, remain}
else
result = {0, capacity, remain - quantity}
redis.call('hincrby', key, 'remain', -quantity)
end

return result

运行步骤

将上面的代码保存到rate_limiter.lua文件。

1
$ vi rate_limiter.lua

将lua脚本文件加载到redis中,返回一个sha值。

1
2
$ redis-cli script load "$(cat rate_limiter.lua)"
"ea8de8016cf04dce431aa9973b8ffe515e06f42a"

在redis客户端中执行evalsha命令,传入步骤2生成的sha值。

1
2
3
4
redis> evalsha ea8de8016cf04dce431aa9973b8ffe515e06f42a 1 ratelimiter 100 30 60
1) (integer) 0 # 未限流
2) (integer) 100 # 容量100
3) (integer) 99 # 当前剩余99

连续执行几次相同的命令,观察效果。

1
2
3
4
5
6
7
8
9
10
11
12
redis> evalsha ea8de8016cf04dce431aa9973b8ffe515e06f42a 1 ratelimiter 10 10 60 5
1) (integer) 0 # 未限流
2) (integer) 10 # 容量10
3) (integer) 5 # 当前剩余5
redis> evalsha ea8de8016cf04dce431aa9973b8ffe515e06f42a 1 ratelimiter 10 10 60 5
1) (integer) 0 # 未限流
2) (integer) 10 # 容量10
3) (integer) 0 # 当前剩余0
redis> evalsha ea8de8016cf04dce431aa9973b8ffe515e06f42a 1 ratelimiter 10 10 60 5
1) (integer) 1 # 已限流
2) (integer) 10 # 容量10
3) (integer) 0 # 当前剩余0

参考链接