Redis + Lua 实现限流（令牌桶算法）

最新推荐文章于 2025-04-09 23:58:10 发布

原创最新推荐文章于 2025-04-09 23:58:10 发布

· 829 阅读

14 ·

版权

文章标签：

#redis #lua #算法

Redis + Lua 可以高效实现分布式限流，确保原子性，避免多个请求同时修改计数器造成不准确的问题。

1️⃣ 基于令牌桶算法的 Redis 限流

基本思路

设定一个 桶（bucket），桶内可以存放一定数量的令牌。
每个请求先从桶中取一个令牌，如果有令牌，则允许请求，否则拒绝请求。
令牌按照一定速率 自动添加 到桶中（但不能超过桶的最大容量）。

2️⃣ 具体实现

实现方式

Redis 存储令牌数量
Lua 以原子性操作检查 & 更新令牌数量

🚀 Lua 脚本

local key = KEYS[1] -- Redis Key（存储令牌数）
local rate = tonumber(ARGV[1]) -- 令牌添加速率（每秒增加多少个）
local capacity = tonumber(ARGV[2]) -- 令牌桶的最大容量
local now = tonumber(ARGV[3]) -- 当前时间戳（秒）
local requested = tonumber(ARGV[4]) -- 请求的令牌数

-- 获取上次请求时的令牌信息
local last_tokens = tonumber(redis.call("GET", key)) or capacity -- 令牌数量（默认满桶）
local last_time = tonumber(redis.call("GET", key .. ":time")) or now -- 上次刷新时间

-- 计算经过的时间
local elapsed = now - last_time

-- 计算新令牌数量
local new_tokens = math.min(capacity, last_tokens + (elapsed * rate))

-- 判断是否允许请求
if new_tokens >= requested then
-- 允许请求，扣除令牌
redis.call("SET", key, new_tokens - requested)
redis.call("SET", key .. ":time", now)
return 1 -- 允许访问
else
-- 令牌不足，拒绝请求
return 0
end

3️⃣ Java 代码调用 Lua

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.util.Pool;
import java.time.Instant;
import java.util.Collections;

public class RedisRateLimiter {
private static final String LUA_SCRIPT =
"local key = KEYS[1] " +
"local rate = tonumber(ARGV[1]) " +
"local capacity = tonumber(ARGV[2]) " +
"local now = tonumber(ARGV[3]) " +
"local requested = tonumber(ARGV[4]) " +
"local last_tokens = tonumber(redis.call('GET', key)) or capacity " +
"local last_time = tonumber(redis.call('GET', key .. ':time')) or now " +
"local elapsed = now - last_time " +
"local new_tokens = math.min(capacity, last_tokens + (elapsed * rate)) " +
"if new_tokens >= requested then " +
" redis.call('SET', key, new_tokens - requested) " +
" redis.call('SET', key .. ':time', now) " +
" return 1 " +
"else " +
" return 0 " +
"end";

private final Pool<Jedis> jedisPool;

public RedisRateLimiter(Pool<Jedis> jedisPool) {
this.jedisPool = jedisPool;
}

public boolean tryAcquire(String key, int rate, int capacity, int requested) {
try (Jedis jedis = jedisPool.getResource()) {
long now = Instant.now().getEpochSecond();
Object result = jedis.eval(LUA_SCRIPT,
Collections.singletonList(key),
Arrays.asList(String.valueOf(rate),
String.valueOf(capacity),
String.valueOf(now),
String.valueOf(requested)));
return Integer.parseInt(result.toString()) == 1;
}
}

public static void main(String[] args) {
JedisPool pool = new JedisPool("localhost", 6379);
RedisRateLimiter limiter = new RedisRateLimiter(pool);

// 测试限流：每秒最多 5 个请求，最大容量 10，当前请求 1 个
for (int i = 0; i < 20; i++) {
boolean allowed = limiter.tryAcquire("rate_limit", 5, 10, 1);
System.out.println("请求 " + (i + 1) + (allowed ? " ✅ 通过" : " ❌ 拒绝"));
try {
Thread.sleep(200); // 200ms 发起一个请求
} catch (InterruptedException ignored) {}
}

pool.close();
}
}