使用Lua Script实现不同的限流算法

Redis中执行Lua Script

redis-cli --eval /tmp/script.lua mykey somekey , arg1 arg2

特别注意：key和arg之间是空格+逗号+空格，否则脚本调用redis-cli命令时会报错

关于Redis CLI的使用建议通读Redis CLI文档，里面有很多常用的操作要熟记，例如如何通过redis-cli连接，如何带密码连接，如果指定lua脚本等

Redis中Debug Lua Script

在Redis中可以使用redis-cli作为debugLua脚本的工具

redis-cli --ldb --eval /tmp/script.lua mykey somekey , arg1 arg2

关于如何调试，官方文档说的也很清楚

更多细节参考Redis官方文档关于Debug Lua一节

以1分钟允许通过最多3个请求为例，分别实现四种算法。

固定窗口

例如在10:01:00 - 10:02:00，这1分钟为例
key为userId:当前窗口的起点，例如 8848:01

local key = KEYS[1]
-- 如果这个key此前不存在则返回-1
-- 关于 or 的写法参考下面图片的截图
local requests = tonumber(redis.call('GET', key) or '-1')
-- 在固定时间范围内允许的最大请求数，例如这里应该是3
local max_requests = tonumber(ARGV[1])
-- 通常是固定窗口的大小,例如60s
local expiry = tonumber(ARGV[2])
-- 当该窗口的key不存在或者未达到最大请求时
if (requests == -1) or (requests < max_requests) then
  -- 自增
  redis.call('INCR', key)
  -- 重新设置该key的过期日期为60s后
  redis.call('EXPIRE', key, expiry)
  return false
else
  return true
end

代码关键部分解释

假如当前时间是10:01:30第一次请求，由于8848:01不存在，所以此时给该键+1并且设置过期时间为60s后。
这里有的同学可能对最后的EXPIRE那里有点争议，其实只是实现方式不同而已。可能大家想到的一种方式是在确定是第一次请求的时候，将剩余的时间给计算出来，等到后面的请求过来的时候，只需要自增即可。
上面这种实现方式就是简化了一下代码，即使到了10:02，因为key变了，所以获取不到10:01的key，这个key自然会在最后一次被设置完过期时间并达到60s后删除

验证

可以看到当3次请求限制达到了，则返回的是true，true被转为1

滑动窗口

固定窗口的问题是：当在10:01:59请求了2次，在10:02:01也请求了2次，这就导致在短短2秒内请求了4次，已经超过了3次的限制。按照固定窗口的逻辑来判断，这两次在各自的时间窗口内是合理的，但是在这种边界时，是不正确的。所以滑动窗口就来解决这个问题

local key = KEYS[1]
-- 当前时间戳
local current_time = tonumber(ARGV[1])
-- 窗口大小，本例中是60 * 1000
local window_size = tonumber(ARGV[3])
-- 本例中是3
local max_requests = tonumber(ARGV[4])
-- 根据当前时间毫秒数 - 超时毫秒数，得到过期时间 expired
local has_expired = current_time - window_size
-- 清除过期的数据
redis.call('ZREMRANGEBYSCORE', key, 0, has_expired)
-- 获取 zset 中的当前元素个数
local current_num = tonumber(redis.call('ZCARD', key))
local next = current_num + 1
-- 达到限流大小 返回 0
if next > max_requests then
  return 0;
else
  -- 往 zset 中添加一个值、得分均为当前时间戳的元素，[value,score]
  redis.call("ZADD", key, current_time, current_time)
  -- 每次访问均重新设置 zset 的过期时间
  redis.call("PEXPIRE", key, window_size)
  return next
end

滑动窗口其实在实际中大部分业务场景下都可以使用，那么为什么还会有令牌桶和漏桶算法的需求呢？带着这个疑问，我在网上找了一些资料，我觉得比较靠谱的回答是流量整型

令牌桶

在实际开发中，令牌桶适用的场景是对请求我们服务器的上游（这里的上游包括来自浏览器的直接请求或者其他不属于我们系统的未知的请求）进行限制。限制请求速率以确保我们的服务器不会被打垮

令牌桶涉及到两个脚本：

• 一个是按照一定的速率往令牌桶里加令牌
• 另外一个是从桶里拿令牌，如果拿到则请求通过，如果拿不到则拒绝请求

添加令牌

-- 当前时间戳
local ts = tonumber(ARGV[1])
-- 设置窗口大小为1s
local min = ts -1
-- 可以为多个key设置添加令牌的速率
for i,key in pairs(KEYS) do
    -- 移除过期的的令牌
    redis.call('ZREMRANGEBYSCORE', key, '-inf', min)
    redis.call('ZADD', key, ts, ts)
    redis.call('EXPIRE', key, 10)
end

获取令牌

-- 当前时间戳
local ts  = tonumber(ARGV[1])
local key = KEYS[1]
local min = ts -1
-- 移除过期的的令牌
redis.call('ZREMRANGEBYSCORE', key, '-inf', min)
-- 只需要计算令牌桶的大小即可，如果令牌桶是有值的，则放行
return redis.call('ZCARD', key)

漏桶

漏桶的适用场景是按照一定的速率向下游服务请求，避免下游服务处理不过来而报错。如果在一段时间内请求数大于固定速率的请求数，则需要排队等待

local ts  = tonumber(ARGV[1])
-- calls per seconds,例如每秒4次,也就是250ms内允许发起1次请求
local cps = tonumber(ARGV[2])
local key = KEYS[1]
local min = ts -1
redis.call('ZREMRANGEBYSCORE', key, '-inf', min)
local last = redis.call('ZRANGE', key, -1, -1)
local next = ts
if type(last) == 'table' and #last > 0 then
  for key,value in pairs(last) do
    -- 最后一个元素 + 固定速率的时间
    next = tonumber(value) + 1/cps
    break
  end
end
if ts > next then
  -- the current ts is > than last+1/cps
  -- we'll keep ts
  next = ts
end
-- 如果ts < next，这里next还是现有zset的最后一个元素 + 1/cps 后的时间
redis.call('ZADD', key, next, next)
-- 必须等待的时间，
-- 如果是ts > next，则这里直接返回0，意味着不等待，直接调用
-- 如果是ts < next, 则说明下一个速率还没达到，则需要等到， next - ts这么长时间
return tostring(next - ts)

参考资料

• 滑动窗口的实现
• 令牌桶和漏桶的实现