用java配合redis 在springboot上实现令牌桶算法

令牌桶算法配合 Redis 在 Java 中的应用令牌桶算法是一种常用的限流算法,适用于控制请求的频率,防止系统过载。结合 Redis 使用可以实现高效的分布式限流。

一.、引入依赖首先,需要在 pom.xml 文件中引入 spring-boot-starter-data-redis 依赖,这个依赖提供了与 Redis 交互的客户端和工具类。

XML 复制代码
<dependencies>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-redis</artifactId>
    </dependency>
</dependencies>

二.、配置 Redis 连接 (我随便用了application.properties ,也可以用 application.yml)

XML 复制代码
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

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 脚本。

java 复制代码
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);
        }
    }
}

五、测试

java 复制代码
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";
        }
    }
}
相关推荐
非ban必选7 分钟前
spring-ai-alibaba第四章阿里dashscope集成百度翻译tool
java·人工智能·spring
非ban必选13 分钟前
spring-ai-alibaba第五章阿里dashscope集成mcp远程天气查询tools
java·后端·spring
纪元A梦30 分钟前
Redis最佳实践——首页推荐与商品列表缓存详解
数据库·redis·缓存
遥不可及~~斌36 分钟前
@ComponentScan注解详解:Spring组件扫描的核心机制
java
高林雨露36 分钟前
Java 与 Kotlin 对比示例学习(三)
java·kotlin
极客先躯1 小时前
高级java每日一道面试题-2025年3月22日-微服务篇[Nacos篇]-Nacos的主要功能有哪些?
java·开发语言·微服务
爱喝醋的雷达1 小时前
Spring SpringBoot 细节总结
java·spring boot·spring
coderzpw2 小时前
当模板方法模式遇上工厂模式:一道优雅的烹饪架构设计
java·模板方法模式
直裾2 小时前
Mapreduce初使用
java·mapreduce
悠夏安末3 小时前
intellij Idea 和 dataGrip下载和安装教程
java·ide·intellij-idea