// 使用 Redis + Lua 脚本实现令牌桶
@Component
public class RateLimiter {
@Autowired
private StringRedisTemplate redis;
// Lua 脚本保证原子性
private static final String LUA_SCRIPT = """
local key = KEYS[1]
local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = tonumber(redis.call('get', key) or '0')
if current < limit then
redis.call('INCR', key)
redis.call('EXPIRE', key, window)
return 1
end
return 0
""";
public boolean isAllowed(String key, int limit, int windowSec) {
DefaultRedisScript<Long> script = new DefaultRedisScript<>();
script.setScriptText(LUA_SCRIPT);
script.setResultType(Long.class);
Long result = redis.execute(script,
Collections.singletonList("rate_limit:" + key),
String.valueOf(limit), String.valueOf(windowSec));
return result != null && result == 1;
}
}
// 拦截器中使用
@Override
public boolean preHandle(HttpServletRequest request, ...) {
String clientId = request.getHeader("X-Client-Id");
if (!rateLimiter.isAllowed(clientId, 100, 60)) {
response.setStatus(429); // Too Many Requests
return false;
}
return true;
}