借助Redis撸个API限流器

借助Redis撸个API限流器

在业务中，很多地方都会用到限流器，为了减少自身服务器压力；为了防止接口滥用，限制请求次数等，今天我们尝试使用redis撸一个API限流器

1. 考虑到分布式，我们取消在内存中限流，引入redis。

2. 考虑到滚动窗口需求，我们取消1.计数器模式，使用2.令牌桶法。

1. 计数器

使用redis的 decr 每到一个请求，计数器减少一，减少到0，则拒绝后续访问，直到下一轮计数器开启，这个设计有一个问题，就是无法避免上一轮计数器结束到下一轮开始这个期间，可以接受2个计数器的访问量。

const Redis = require('ioredis');
const redis = new Redis(redisConfig);
const rateLimitKey = 'rateLimitKey', pexpire = 60000, limit = 100, amount = 1;
const ttl = await redis.pttl(rateLimitKey)
if (ttl < 0) {
  await redis.psetex(rateLimitKey, pexpire, limit - amount)
  return {
    limit,
    remain: limit - amount,
    rejected: false,
    retryDelta: 0,
  };
} else {
  const remain = await redis.decrby(rateLimitKey, amount)
  return {
    limit,
    remain: remain > 0 ? remain : 0,
    rejected: remain >= 0,
    retryDelta: remain > 0 ? 0 : ttl
  }
}

2.令牌桶

一边不断消耗令牌，另一边不断流入到桶中。使用其他进程去操作流入令牌到桶，是个错误的想法，随着KEY的增多，这无疑是最大的负载。所以我们考虑通过记录上一次请求时间与剩余令牌，去计算应该流入令牌数量，但是这个需要多次操作redis，考虑到竞争条件，我们选择用lua脚本去做。

• 流入令牌数量计算：math.max(((nowTimeStamp - lastTimeStamp) / pexpire) * limit, 0)

const Redis = require('ioredis');
const client = new Redis(redisConfig)
    client.defineCommand('rateLimit', {
        numberOfKeys: 2,
        lua: fs.readFileSync(path.join(__dirname, './rateLimit.lua'), {encoding: 'utf8'}),
    })
const args = [`${Key}:V`, `${Key}:T`,Limit, Pexpire, amount];
const [limit, remain, rejected, retryDelta] = await client.rateLimit(...args)
return {
        limit,
        remain,
        rejected: Boolean(rejected),
        retryDelta,
    }

lua代码如下: rateLimit.lua

local valueKey = KEYS[1]                    -- 存储计数器的KEY
local timeStampKey = KEYS[2]                -- 存储上次访问时间戳的KEY
local limit = tonumber(ARGV[1])             -- 单位时间内可访问的次数
local pexpire = tonumber(ARGV[2])           -- KEY的有效期 ms
local amount = tonumber(ARGV[3])            -- 每次减少次数
redis.replicate_commands()
local time = redis.call('TIME')
local nowTimeStamp = math.floor((time[1] * 1000) + (time[2] / 1000))
local nowValue
local lastValue = redis.call('GET', valueKey)     -- 计数器剩余量
local lastTimeStamp                               -- 上次更新时间
if lastValue == false then
    lastValue = 0
    lastTimeStamp = nowTimeStamp - pexpire
else
    lastTimeStamp = redis.call('GET', timeStampKey)
    if(lastTimeStamp == false) then
        lastTimeStamp = nowTimeStamp - ((lastValue / limit) * pexpire)
    end
end
local addValue = math.max(((nowTimeStamp - lastTimeStamp) / pexpire) * limit, 0)     -- 本次增加的计数次数
nowValue = math.min(lastValue + addValue, limit)
local remain = nowValue - amount
local rejected = false
local retryDelta = 0
if remain < 0 then
    remain = 0
    rejected = true
    retryDelta = math.ceil(((amount - nowValue) / limit) * pexpire)
else
    if (remain - amount) < 0 then
        retryDelta = math.ceil((math.abs(remain - amount) / limit) * pexpire)
    end
end
if rejected == false then
    redis.call('PSETEX', valueKey, pexpire, remain)
    if addValue > 0 then
        redis.call('PSETEX', timeStampKey, pexpire, nowTimeStamp)
    else
        redis.call('PEXPIRE', timeStampKey, pexpire)
    end
end
return { limit, remain, rejected, retryDelta }