用java配合redis 在springboot上实现令牌桶算法
令牌桶算法配合 Redis 在 Java 中的应用令牌桶算法是一种常用的限流算法，适用于控制请求的频率，防止系统过载。结合 Redis 使用可以实现高效的分布式限流。
一.、引入依赖首先，需要在 pom.xml 文件中引入 spring-boot-starter-data-redis 依赖，这个依赖提供了与 Redis 交互的客户端和工具类。
<dependencies><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-data-redis</artifactId></dependency>
</dependencies>
二.、配置 Redis 连接 （我随便用了application.properties ，也可以用 application.yml）
spring.redis.host=localhost
spring.redis.port=6379
spring.redis.password=
spring.redis.jedis.pool.max-active=8
spring.redis.jedis.pool.max-wait=-1
spring.redis.jedis.pool.max-idle=8
spring.redis.jedis.pool.min-idle=0
spring.redis.timeout=2000三.、实现令牌桶算法
使用 Lua 脚本和 Redis 操作来实现令牌桶算法。在resources创建一个 Lua 脚本文件 request_rate_limiter.lua
-- 获取到限流资源令牌数的key和响应时间戳的key
local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
-- 分别获取填充速率、令牌桶容量、当前时间戳、消耗令牌数
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
-- 计算出失效时间，大概是令牌桶填满时间的两倍
local fill_time = capacity / rate
local ttl = math.floor(fill_time * 2)
-- 获取到最近一次的剩余令牌数，如果不存在说明令牌桶是满的
local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil thenlast_tokens = capacity
end
-- 上次消耗令牌的时间戳，不存在视为0
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil thenlast_refreshed = 0
end
-- 计算出间隔时间
local delta = math.max(0, now - last_refreshed)
-- 剩余令牌数量 = “令牌桶容量” 和 “最后令牌数+（填充速率*时间间隔）”之间的最小值
local filled_tokens = math.min(capacity, last_tokens + (delta * rate))
-- 如果剩余令牌数量大于等于消耗令牌的数量则流量通过，否则不通过
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed thennew_tokens = filled_tokens - requestedallowed_num = 1
end
-- 更新令牌桶状态
if ttl > 0 thenredis.call("setex", tokens_key, ttl, new_tokens)redis.call("setex", timestamp_key, ttl, now)
end
return allowed_num
四、创建一个 TokenBucket 类，使用 StringRedisTemplate 执行 Lua 脚本。
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;import javax.xml.crypto.Data;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;@Component
public class TokenBucket {private static final String TOKEN_BUCKET_KEY_PREFIX = "rate_limiter:";private static final String LUA_SCRIPT_PATH = "request_rate_limiter.lua";@Autowiredprivate StringRedisTemplate stringRedisTemplate;public boolean tryAccess(String key, int limitCount, int refillRate) {String luaScript = loadLuaScript();// 使用加载的Lua脚本创建一个RedisScript 对象。 DefaultRedisScript 是Spring Data Redis提供的一个类，用于封装Lua脚本。RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);// TOKEN_BUCKET_KEY_PREFIX + key + ":tokens"：用于存储令牌桶中的令牌数量。•// TOKEN_BUCKET_KEY_PREFIX + key + ":timestamp"：用于存储上一次令牌桶被填充的时间戳。List<String> keys = Arrays.asList(TOKEN_BUCKET_KEY_PREFIX + key + ":tokens", TOKEN_BUCKET_KEY_PREFIX + key + ":timestamp");//执行Lua脚本/*** 使用 StringRedisTemplate 执行Lua脚本。execute方法的参数包括：* • redisScript ：Lua脚本对象。* • keys：Redis键列表。* • String.valueOf(refillRate)：令牌桶的填充速率。* • String.valueOf(limitCount)：令牌桶的最大容量。* • String.valueOf(System.currentTimeMillis() / 1000)：当前时间戳（秒级）。* • "1" ：请求的令牌数量（这里假设每次请求需要1个令牌）。*/Long result = stringRedisTemplate.execute(redisScript, //Lua脚本对象。keys, //Redis键列表String.valueOf(refillRate), //令牌桶的填充速率String.valueOf(limitCount), //令牌桶的最大容量String.valueOf(System.currentTimeMillis() / 1000), //当前时间戳（秒级）"1" //请求的令牌数量（这里假设每次请求需要1个令牌）。);return result != null && result == 1;}private String loadLuaScript() {InputStreamReader reader = null;try {//使用类加载器从指定路径 LUA_SCRIPT_PATH 获取资源流。我的lua文件在resources根目录下面//使用 InputStreamReader 将输入流 resourceStream 包装成一个字符流，并指定字符编码为UTF-8。这样可以确保读取的文件内容是正确的编码格式reader = new InputStreamReader(getClass().getClassLoader().getResourceAsStream(LUA_SCRIPT_PATH), StandardCharsets.UTF_8);/*创建一个Scanner对象，用于读取字符流。useDelimiter("\\A")：设置分隔符为文件的开始标记 \\A 。这意味着Scanner会将整个文件内容视为一个单一的字符串。next() ：读取并返回整个文件内容*/return new java.util.Scanner(reader).useDelimiter("\\A").next();} catch (Exception e) {throw new RuntimeException("Failed to load Lua script", e);}}
}
五、测试
package com.lhx.testany.controller;import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RestController;@RestController
public class RateLimiterController {@Autowiredprivate TokenBucket tokenBucket;@GetMapping("/test")public String testRateLimiter() {//为了测试，把最多设成2个令牌，每秒只生成1个令牌, 我快速在浏览器调用时，//生成的令牌，不足于被消耗，就会执行esle,这样就能做限流if (tokenBucket.tryAccess("api1", 2, 1)) {return "Request allowed";} else {return "Request denied due to rate limit";}}
}