使用redis来实现分布式限流的正确姿势

cpp可以是hiredis作为redis client，其他语言则更方便。

由于redis 指令操作逻辑都跑在一个主线程上（包括lua脚本），因此使用redis可以原子的执行lua脚本指令，即使redis有多个分片。不过在多分片情况下lua脚本的KEY必须能够被hash到同一个分片才可以，如果业务的KEY本身不能够保证被hash到同一个slot的话则可以使用redis本身提供的{}操作符来指定shard key。

限流逻辑的关键就在于LUA脚本怎么写。这里主要实现了滑动时间窗口限流和令牌桶限流两种比较常用的分布式限流方法。下面上代码

令牌桶：

-- 获取到限流资源令牌数的key和响应时间戳的key
local last_left_tokens_key = KEYS[1];
local last_check_time_key = KEYS[2];
        -- 填充速率
local token_fill_rate = tonumber(ARGV[1]);
        -- 令牌桶容量
local bucket_capacity = tonumber(ARGV[2]);
        -- 当前时间戳
local now = tonumber(ARGV[3]);
        -- 请求的令牌数
local requested_tokens = tonumber(ARGV[4]);
        -- 计算下key得失效时间，防止已经不用的限流key长期占用内存
local fill_up_time = bucket_capacity / token_fill_rate;
local ttl = math.floor(fill_up_time * 2);
        -- 获取到最近一次的剩余令牌数，如果不存在说明令牌桶是满的
local last_left_tokens = tonumber(redis.call("GET", last_left_tokens_key));
if last_left_tokens == nil then
	last_left_tokens = bucket_capacity;
end
        -- 获取下上次消耗令牌的时间戳
local last_check_time = tonumber(redis.call("GET", last_check_time_key));
if last_check_time == nil then
	last_check_time = 0;
end
        -- 计算下前后两次得时间间隔
local interval = math.max(0, now - last_check_time);
        -- 计算当前剩余的token数量
local cur_tokens = 0;
if (last_check_time ~= now) then
	cur_tokens = math.min(bucket_capacity, last_left_tokens + math.floor(interval * token_fill_rate));
else
	cur_tokens = math.min(bucket_capacity, last_left_tokens);
end
local check_flag = 0;
if cur_tokens < requested_tokens then
	check_flag = 1;
end
local after_consuming_tokens = cur_tokens;
if check_flag == 0 then
	after_consuming_tokens = cur_tokens - requested_tokens;
end
if (after_consuming_tokens ~= last_left_tokens or check_flag == 0) then
	redis.call("SETEX", last_left_tokens_key, ttl, after_consuming_tokens);
	redis.call("SETEX", last_check_time_key, ttl, now);
end
return check_flag;

滑动时间窗口：

local time_window_key = KEYS[1];
        -- 把时间窗口左边的数据都移除
redis.call('ZREMRANGEBYSCORE', time_window_key, 0, ARGV[1]);
        -- 移除之后获取下当前的数量
local cnt = redis.call('ZCARD', time_window_key);
        -- 判断是否超过了限制
if ((cnt == nil) or (cnt < tonumber(ARGV[3]))) then
	redis.call('ZADD', time_window_key, tonumber(ARGV[2]), ARGV[4]);
	return 0;
else
	return 1;
end