分布式限流——基于Redis的Lua脚本限流实现

发布时间 2024-01-13 18:33:57作者: 想进美团

分布式限流

当你的应用分布式部署出现对等端(peer)时,单机的限流往往不能满足对下游保护的作用,因为它仅仅是jvm内存层面的流量控制。这个时候自然而然会想到用一些跨JVM的分布式中间件控制在单位时间窗口内的请求是否通行,本文我们将探讨如何借助Redis实现分布式限流。

1 固定窗口限流

前文已经介绍了固定时间窗口的限流原理和算法,此处不再赘述。直接上代码:

@Slf4j
@Component
public class DistributedFixedRateLimiter {

    @Autowired
    private RedisTemplate<Object, Object> redisTemplate;

    public boolean tryAcquire( String key,int windowSize, int maxRequestCount){
        String script = getFixedRateLimiterLuaScript();
        Long res = redisTemplate.execute(new DefaultRedisScript<>(script, Long.class),
                Collections.singletonList(key),
                maxRequestCount, windowSize);
        if (res != null && res == -1) {
            log.info("请求失败");

            return true;
        }else {
            log.info("--- 请求成功 , 剩余可用请求数:{}", res);
            return false;
        }
    }
}

这里的getFixedRateLimiterLuaScript()方法就是获取Lua脚本,Lua脚本代码如下:

local key = KEYS[1] -- 限流资源
local maxRequestCount = ARGV[1] -- 限流请求数
local windowSize = ARGV[2] -- 限流时间
local currentCount = redis.call('get', key) -- 当前请求数
-- 限流存在并且超过限流大小,则返回剩余可用请求数=0
if (currentCount and tonumber(currentCount) >= tonumber(maxRequestCount)) then
    return -1
end
-- 请求数自增
currentCount = redis.call('incr', key)
-- 第一次请求,则设置过期时间
if (tonumber(currentCount) == 1) then
    redis.call('expire', key, windowSize)
end
-- 返回剩余可用请求数
return tonumber(maxRequestCount) - tonumber(currentCount)

测试:

@Test
public void test05() {
    String key = "/fixed/window";
    for (int i = 0; i < 10; i++) {
        boolean b = limiter.tryAcquire(key,60, 5);
    }
}

image-20240112222838911

2 滑动窗口限流

基于redis+lua脚本的滑动窗口限流和前文中介绍的略有不同,并不是通过将时间窗口分割成固定大小的子窗口,以一个子窗口为单位移动,而是根据每次请求的间隔,释放掉这次请求的时间减去时间窗口之前的请求,以此来移动。

具体的,通过zset实现滑动窗口,其中zset的key为请求资源,member为单次的请求id,score为当前时间戳:

@Slf4j
@Component
public class DistributedSlideWindowRateLimiter {

    @Resource
    private RedisTemplate<Object, Object> redisTemplate;

    public boolean tryAcquire(String key, String currentTimeKey ,long windowSize, int limitCount, long currentTime) throws Exception {
        String script = getSlideWindowLuaScript();
        List<Long> result = redisTemplate.execute(new DefaultRedisScript<>(script, List.class),
                Arrays.asList(new String[]{key, currentTimeKey}),
                windowSize, limitCount, currentTime);
        if (result == null)
            throw new Exception("redis execute occur exception!");
        log.info("剩余请求数:{}", result.get(1));
        return result.get(0) == 1;
    }

    private String getSlideWindowLuaScript() {
        return "local key = KEYS[1]  -- 限流关键字\n" +
                "local current_time_key = KEYS[2]   -- 当前时间戳的key\n"+
                "local window_size = tonumber(ARGV[1])  -- 滑动窗口大小\n" +
                "local limit = tonumber(ARGV[2])  -- 限制的请求数\n" +
                "local current_time = tonumber(ARGV[3])  -- 当前时间戳\n" +
                "local last_requested = 0   -- 已经用掉的请求数\n"+
                "local remain_request = 0   -- 剩余可以分配的请求数\n"+
                "local allowed_num = 0  -- 本次允许通过的请求数\n" +
                "\n" +
                "local exists_key = redis.call('exists', key)\n" +
                "if (exists_key == 1) then\n" +
                "    last_requested = redis.call('zcard', key)\n" +
                "end\n"+
                "remain_request = limit - last_requested\n" +
                "if (last_requested < limit) then\n" +
                "    allowed_num = 1\n" +
                "    redis.call('zadd', key, current_time, current_time_key)\n" +
                "end\n" +
                "redis.call('zremrangebyscore', key, 0, current_time - window_size)\n" +
                "redis.call('expire', key, window_size)\n" +
                "\n" +
                "return { allowed_num, remain_request }";
    }
}

需要注意,返回的remain_request是请求前剩余的可用请求数。

为了更直观的看出效果,测试时每隔1秒请求一次,通过日志输出可以看出,每隔1秒请求,每次请求可以淘汰5秒前的请求,所以基本上都会有1个请求可用:

2024-01-13 14:39:10.783  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:5
2024-01-13 14:39:11.797  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:11.800  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:4
2024-01-13 14:39:12.808  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:12.810  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:3
2024-01-13 14:39:13.813  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:13.816  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:2
2024-01-13 14:39:14.816  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:14.818  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:15.821  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:15.823  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:0
2024-01-13 14:39:16.828  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求失败
2024-01-13 14:39:16.830  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:17.832  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:17.834  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:18.835  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:18.837  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:19.837  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:19.840  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:20.841  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:20.843  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:21.851  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:21.853  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:0
2024-01-13 14:39:22.856  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求失败
2024-01-13 14:39:22.857  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:23.863  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:23.865  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:24.870  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:24.872  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:25.877  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:25.880  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:26.893  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:26.895  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:27.898  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:27.900  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:0
2024-01-13 14:39:28.916  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求失败
2024-01-13 14:39:28.918  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:29.924  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功
2024-01-13 14:39:29.925  INFO 7296 --- [           main] .r.d.s.DistributedSlideWindowRateLimiter : 剩余请求数:1
2024-01-13 14:39:30.930  INFO 7296 --- [           main] com.dianping.ratelimiter.AppTest         : 请求成功

通过redis中zset成员的变化也不难看出,每次请求会淘汰5秒前的请求:

image-20240113144028957

image-20240113144037274

3 令牌桶算法

基于redis的令牌桶算法可以用hash实现,其中hash的key为请求资源,lua中内设字段tokens_count_fieldlast_refreshed_field分别表示桶中的令牌数和上次更新时间,编写Lua脚本时需要注意一点,rate速率在Java代码中的语义是每秒生成令牌的数量,在Lua中计算的时间差是用时间戳计算的,单位是毫秒,所以需要转换为秒。具体如下:

local key = KEYS[1]
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local current_time = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
local tokens_count_field = 'bucket_token_count'
local last_refreshed_field = 'last_update_time'

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)

-- 如果过期了,先让bucket充满令牌
local last_tokens = tonumber(redis.call("hget", key, tokens_count_field))
if last_tokens == nil then
  last_tokens = capacity
end

local last_refreshed = tonumber(redis.call("hget", key, last_refreshed_field))
if last_refreshed == nil then
  last_refreshed = 0
end

local delta = math.max(0, current_time - last_refreshed)
local filled_tokens = math.min(capacity, math.floor(last_tokens + (delta / 1000 * rate)))
local allowed = filled_tokens >= requested
local remain_tokens = filled_tokens
local allowed_num = 0
if allowed then
  remain_tokens = filled_tokens - requested
  allowed_num = 1
end

redis.call("hset", key, tokens_count_field, remain_tokens)
redis.call("hset", key, last_refreshed_field, current_time)
redis.call("expire", key, ttl)

return { allowed_num, remain_tokens }

在测试代码时,我发现在计算filled_tokens时,理论上来说令牌数应该是整型,所以写代码时我下意识加了math.floor()向下取整,但是时间戳往往是很零碎的值(因为以毫秒为单位),如果每次保存时间间隔内生成的令牌时都向下取整,将会“损失精度”,使生成令牌的速率更”离散化“,而这样在突发流量到来时,由于请求间隔小(每次都是趋于0的状态),导致每次生成的令牌数都是0,通过日志输出可以证明这一猜想:

测试代码:(每秒生成1个令牌,桶容量为10,每次请求需要3个令牌)

@Test
public void test07() throws InterruptedException {
    String key = "/token/bucket";
    for (int i = 0; i < 20; i++) {
        long currentTime = System.currentTimeMillis();
        Thread.sleep(500);
        boolean b = tokenBucketRateLimiter.tryAcquire(key, 1, 10, currentTime, 3);
    }
}

以下是向下取整的输出结果:

2024-01-13 18:20:27.959  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:7
2024-01-13 18:20:28.474  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:5
2024-01-13 18:20:28.977  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:2
2024-01-13 18:20:29.488  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:30.004  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:30.509  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:31.012  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:31.518  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:32.022  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:32.524  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:33.039  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:33.546  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:34.051  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:34.566  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:35.071  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:35.575  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:36.082  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:36.584  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:37.087  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:20:37.595  INFO 1368 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败

以下是不取整的输出结果:

2024-01-13 18:14:00.753  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:7
2024-01-13 18:14:01.256  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:5
2024-01-13 18:14:01.759  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:2
2024-01-13 18:14:02.265  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:0
2024-01-13 18:14:02.785  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:03.289  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:03.792  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:04.295  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:04.801  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:05.307  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:0
2024-01-13 18:14:05.813  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:06.318  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:06.824  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:07.327  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:07.833  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:08.338  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求成功, 剩余令牌:0
2024-01-13 18:14:08.844  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:09.351  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:09.857  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败
2024-01-13 18:14:10.360  INFO 14144 --- [           main] .r.d.t.DistributedTokenBucketRateLimiter : 请求失败

image-20240113175043019

所以经过对比,我认为还是不应该取整,因为在如果生成的速率小于消费的速率,即短时间内到来的大量请求将会被阻塞。

本博客内容仅供个人学习使用,禁止用于商业用途。转载需注明出处并链接至原文。