分布式限流——基于Redis的Lua脚本限流实现
分布式限流
当你的应用分布式部署出现对等端（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);
}
}
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秒前的请求：
3 令牌桶算法
基于redis的令牌桶算法可以用hash实现，其中hash的key为请求资源，lua中内设字段tokens_count_field和last_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 : 请求失败
所以经过对比，我认为还是不应该取整，因为在如果生成的速率小于消费的速率，即短时间内到来的大量请求将会被阻塞。
完整代码已上传GitHub：https://github.com/zhaobo97/RateLimiter
本博客内容仅供个人学习使用，禁止用于商业用途。转载需注明出处并链接至原文。
