redis中使用lua脚本及使用脚本完成限流

vlambda
2020-07-11

redis中使用lua脚本及使用脚本完成限流

redis支持lua脚本，可以在lua脚本中将多个redis执行单元组合在一起，完成原子性操作。先来看一个使用lua的简单示例：

eval "if redis.call('get',KEYS[1]) == ARGV[1] then return 0 else return -1 end" 1 name star

基本语法就是这样，KEYS用来传递redis要使用key值，ARGV用来传递脚本需要的值参数（你要倒行逆施用ARGV传key值，KEYS当作value也是ok的），只是最后要指定KEYS参数个数，紧跟着KEYS参数，最后跟着ARGV参数。

搞清楚了基础语法，就可以使用redis命令组合来完成一个限流操作了，类似数据库存储过程，在jedis中使用和在redis-cli中使用是一致的。

我们先写一个简单计数器限流，为指定的限流key设置过期时间（单位限流时间），超过限流值时，返回失败。

local curr_count = tonumber(redis.call('incr',KEYS[1]))
if curr_count == 1 then
	redis.call('pexpire',KEYS[1],tonumber(ARGV[1]))
elseif curr_count > tonumber(ARGV[2]) then
	return -1
end
return curr_count

可以将脚本加载到redis，通过sha值调用脚本。

src/redis-cli -h 127.0.0.1 -a redis5 script load "$(cat lua-scripts/count_limit.lua)"
#得到sha值
evalsha $sha  1  day  2000  1

2000毫秒内，限流一个，可以执行看到效果，两秒内多次执行返回-1。实现起来简单，但限流不够精准，会出现很多极端情况。比如两个时间窗口临界处出现2倍的请求，第一个2秒内最后一毫秒出现一个请求，第二个2秒内的第一个毫秒出现一个请求等等。

所以时间粒度应尽量细一些，比如每秒1000请求，可以细分为100毫秒100个，限流就是为了保护应用不被拖垮。

下面使用令牌桶算法实现限流，令牌桶其实是类似计数算法的，基本思想就是：先初始化一个满的令牌桶（数量是有上限的）；每个请求来了都要去令牌桶获取令牌（令牌消耗）；令牌桶以一定的间隔周期重新生成令牌放入桶中。

local capacity = tonumber(ARGV[1])
local now = tonumber(ARGV[2])
local rate = tonumber(ARGV[3])

local curr_tokens = redis.call('get',KEYS[1])
local ttl = capacity / rate * 2
local new_tokens = 0
if curr_tokens then
	local last_millis = redis.call('get',KEYS[1]..'_millis')
	local new_add = math.floor((now - last_millis) *  rate)
	new_tokens = math.min(capacity, curr_tokens + new_add) -1
	if new_tokens < 1 then
		return -1
	end
else
	new_tokens = capacity -1
end
redis.call('set',KEYS[1],new_tokens,'PX',ttl)
redis.call('set',KEYS[1]..'_millis',now,'PX',ttl)
return new_tokens