利用redis实现令牌桶限流

lua脚本如下

argument	mean
ARGV[1]	rate 「每秒生成几个令牌」
ARGV[2]	burst 「令牌桶最大值」
ARGV[3]	now_time「当前时间戳」
ARGV[4]	get token nums 「开发者需要获取的token数」
KEYS[1]	表示资源的tokenkey
KEYS[2]	表示刷新时间的key

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)
local last_tokens = tonumber(redis.call("get", KEYS[1]))
if last_tokens == nil then
    last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("get", KEYS[2]))
if last_refreshed == nil then
    last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
if allowed then
    new_tokens = filled_tokens - requested
end

redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)

return allowed

redis 的eval命令是用来执行lua脚本的

命令格式：EVAL script numkeys key [key …] arg [arg …]

• script 参数是一段 Lua5.1 脚本程序。脚本不必(也不应该[^1])定义为一个 Lua 函数
• numkeys 指定后续参数有几个key，即：key [key …]中key的个数。如没有key，则为0
• key [key …] 从 EVAL 的第三个参数开始算起，表示在脚本中所用到的那些 Redis 键(key)。在Lua脚本中通过KEYS[1], KEYS[2]获取。
• arg [arg …] 附加参数。在Lua脚本中通过ARGV[1],ARGV[2]获取。

// 例子：使用redis为lua内置的redis.call函数
//      脚本内容为：先执行SET命令，在执行EXPIRE命令(相当于执行SETEX key1 60 10)
//      numkeys=1，keys数组有一个元素userAge（代表redis的key）
//      arg数组元素中有两个元素：10（代表userAge对应的value）和60（代表redis的存活时间）
127.0.0.1:6379> EVAL "redis.call('SET', KEYS[1], ARGV[1]);redis.call('EXPIRE', KEYS[1], ARGV[2]); return 1;" 1 userAge 10 60
(integer) 1
127.0.0.1:6379> get userAge
"10"
127.0.0.1:6379> ttl userAge
(integer) 50

注意避坑

为什么now需要通过参数传递进去，而不是在脚本内部自己获取？我们先来看个例子

127.0.0.1:6379> eval "local now = redis.call('time')[1]; redis.call('set','now',now); return redis.call('get','now')" 0
(error) ERR Error running script (call to f_e745355f11745192bd45376618a34bec9145653b): @user_script:1: @user_script: 1: Write commands not allowed after non deterministic commands. Call redis.replicate_commands() at the start of your script in order to switch to single commands replication mode.

这是因为在Redis中time命令是一个随机命令（时间是变化的），出于主从数据一致性考虑，要求脚本必须是纯函数的形式，也就是说对于一段Lua脚本给定相同的参数，写入Redis的数据也必须是相同的，对于随机性的写入Redis是拒绝的，所以在Lua脚本中调用了随机命令之后禁止再调用写命令