用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 then
  last_tokens = capacity
end
-- 上次消耗令牌的时间戳，不存在视为0
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_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 then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end
-- 更新令牌桶状态
if ttl > 0 then
  redis.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";

    @Autowired
    private 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 {

    @Autowired
    private  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";
        }
    }
}