lua脚本实现滑动窗口的分布式全局限流器, 控制api接口qps

zzzgd816

1672人浏览 · 2023-04-10 21:02:01

zzzgd816 · 2023-04-10 21:02:01 发布

更新
2023年4月19日
发现坑,测试环境没问题,生产报 Write commands not allowed after non deterministic commands.

原因大概是redis在哨兵环境下主从复制是基于命令重复执行的复制, 也就对time这种不确定结果的命令,考虑数据一致性而不执行 zadd 抛异常
解决办法是加入redis.replicate_commands()

建议大家使用lua脚本, 先不要急着投入使用, 而是留一个后门接口, 发版后调用接口测试无问题后, 下个版本在使用

参考资料 :

前言

限流器, 从算法实现的角度来说, 就我知道的常见的有滴漏桶, 令牌桶, 滑动窗口计数,固定窗口计数法

从实现的工具来说, 常见的有 guava的 RateLimiter (令牌桶)
redis的每秒或者每一分钟过期时间的incr(固定窗口计数)

但是这些大多数时候都被我们用来当做单个机器上的限流措施, 尤其是guava这种单体框架. redis的incr虽然能控制全局, 但是还是有问题.

场景

对接支付系统的接口, 给出的qps是5, 一开始是直接在代码中sleep200毫秒, 后面发现分布式的集群这种限制是行不通的.
使用incr方法? 不行, 因为你的窗口期1s不一定就是下游系统的窗口期. 比如下面这样, 虽然在我们这边保证了1秒内请求5次, 但是也许在别人系统的窗口中正好是你两个窗口期集中请求的部分, 就会报错
[00 000][0000 0]
[…] [111111]

lua脚本

这里采用滑动窗口计数的方式, 这样能保证不管什么时间段, 这个前后1s内总是不会超过5个请求, 也就能满足下游的要求

加入
– 告诉redis以redis命令的形式实现主从复制
redis.replicate_commands() 命令

-- 获取zset的key
local key = KEYS[1]
-- 脚本传入的限流大小
local limit = tonumber(ARGV[1])
-- 区间长度
local rangeMill = tonumber(ARGV[2])
-- 告诉redis以redis命令的形式实现主从复制
redis.replicate_commands()
-- 当前网络时间戳 秒和这秒过去的微秒数
local date =  redis.call("time")
-- 得到微妙时间戳 微秒
local now = tonumber(date[1])*1000000 + tonumber(date[2])
-- 为了避免now这个时间和后面zount时候的偏差,将右边界增加10s
local toEnd = now + 10 * 1000000
-- 根据区间获取起始时间戳 微秒
local start = now - rangeMill*1000
-- 过期时间s,保险起见+1s
local expireSec = tonumber(math.ceil(rangeMill / 1000)) + 1
if expireSec == 0 then
    expireSec = 1
end
-- 脚本传入的uuid
local uuid = ARGV[3]
-- 获取当前流量总数
local count = tonumber(redis.call('zcount',key, start, toEnd))
--是否超出限流值
if count + 1 >limit then
    return false
-- 不需要限流
else
    -- 添加当前访问时间戳到zset
    redis.call('zadd', key, now, uuid)
    -- 移除时间区间以外不用的数据，不然会导致zset过大,start前1s的数据
    redis.call('zremrangebyscore',key, 0, start-1000000)
    redis.call("expire", key,expireSec)
    return true
end

然后使用redisTemplate去调用这个脚本

    /**
     * 读取lua脚本
     * @return
     */
    @Bean("limitRateScript")
    public RedisScript<Boolean> loadRedisScript(){
        DefaultRedisScript<Boolean> redisScript = new DefaultRedisScript<>();
        //lua脚本路径
        redisScript.setLocation(new ClassPathResource("lua/token_rate_limit.lua"));
//        redisScript.setScriptText(lua);
        //lua脚本返回值
        redisScript.setResultType(java.lang.Boolean.class);
        return redisScript;
    }

/**
     * 通过lua脚本实现令牌桶
     *
     * @param key        key
     * @param limit      限制次数
     * @param rangeMills 区间,ms
     * @return
     */
    public static boolean limit(String key, int limit, int rangeMills) {
        RedisScript<Boolean> redisScript = SpringCtxUtil.getBean("limitRateScript", RedisScript.class);
        //当前时间戳
        return Boolean.TRUE.equals(getStringRedisTemplate().execute(
                //lua限流脚本
                redisScript,
                //限流资源名称
                Collections.singletonList(key),
                //限流大小
                String.valueOf(limit),
                //限流窗口的区间大小
                String.valueOf(rangeMills),
                //id值，保证zset集合里面不重复，不然会覆盖
                UUID.randomUUID().toString()

        ));
    }

    /**
     * 使用令牌桶限流,但是有一个等待时间
     *
     * @param key         key
     * @param limit       限制次数
     * @param rangeMill   区间,ms
     * @param waitTimeOut
     * @param timeUnit
     */
    public static void limitWait(String key, int limit, int rangeMill, int waitTimeOut, TimeUnit timeUnit) {
        long startTm = System.currentTimeMillis();
        while (System.currentTimeMillis() - startTm <= timeUnit.toMillis((long) waitTimeOut)) {
            boolean b = limit(key, limit, rangeMill);
            if (b) {
                return;
            }
            try {
                TimeUnit.MILLISECONDS.sleep(100L);
            } catch (InterruptedException var9) {
                var9.printStackTrace();
            }
        }
        throw new Exception("wait limit lock timeout|" + key);
    }

实际效果

使用jmeter模拟100线程循环5次
大概27s执行完, 其中失败率77%, 即成功的是 500 * 23%=115
符合qps=5的预期 (比27*5少, 其中少的部分是因为基本不可能做好所有请求都等上一个窗口完了以后完美衔接)

附

如果把窗口边界时间作为参数给lua, 会出现超频的问题
最开始, 窗口的起始时间, 结束时间, 都是由业务传给lua的, 但是通过jemter发现还是触发了qps超频的错误,尤其是启动后第一次批量请求

原因大概是因为一开始大量的请求挤进来, 每个线程都生成相近的边界时间, 比如A是 09:00:00 000 ~ 09:00:01 000,
B是 09:00:00 100 ~ 09:00:01 100, 但是B却优先从jedis连接池拿到了redis连接, 然后在09:00:01 100处记数+1, A再去判断的时候就会漏掉B这个数, 导致A也通过