Java基于Redis的滚动限流与回退机制实现

<dependency>
    <groupId>com.giffing.bucket4j.core</groupId>
    <artifactId>bucket4j-core</artifactId>
    <version>8.1.1</version> <!-- 使用最新稳定版本 -->
</dependency>
<dependency>
    <groupId>com.giffing.bucket4j.redis</groupId>
    <artifactId>bucket4j-redis</artifactId>
    <version>8.1.1</version> <!-- 确保与core版本一致 -->
</dependency>
<dependency>
    <groupId>redis.clients</groupId>
    <artifactId>jedis</artifactId>
    <version>4.3.1</version> <!-- 或使用Lettuce等其他Redis客户端 -->
</dependency>

import io.github.bucket4j.Bandwidth;
import io.github.bucket4j.Bucket;
import io.github.bucket4j.Bucket4j;
import io.github.bucket4j.ConsumptionProbe;
import io.github.bucket4j.redis.jedis.Bucket4jJedis;
import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;

import java.time.Duration;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

public class RedisRateLimiter {

    private final JedisPool jedisPool;
    private final ConcurrentHashMap<String, Bucket> buckets = new ConcurrentHashMap<>();
    private final long capacity; // 令牌桶容量
    private final long refillTokens; // 每次补充的令牌数
    private final Duration refillInterval; // 补充令牌的时间间隔

    /**
     * 构造函数
     * @param redisHost Redis主机地址
     * @param redisPort Redis端口
     * @param capacity 令牌桶容量
     * @param refillTokens 每次补充的令牌数
     * @param refillInterval 补充令牌的时间间隔
     */
    public RedisRateLimiter(String redisHost, int redisPort, long capacity, long refillTokens, Duration refillInterval) {
        JedisPoolConfig poolConfig = new JedisPoolConfig();
        poolConfig.setMaxTotal(100); // 最大连接数
        poolConfig.setMaxIdle(20);   // 最大空闲连接数
        poolConfig.setMinIdle(5);    // 最小空闲连接数
        this.jedisPool = new JedisPool(poolConfig, redisHost, redisPort);
        this.capacity = capacity;
        this.refillTokens = refillTokens;
        this.refillInterval = refillInterval;
    }

    /**
     * 获取或创建指定键的限流桶
     * @param key 限流键 (例如: 用户ID, IP地址, API路径)
     * @return 对应的Bucket实例
     */
    public Bucket getOrCreateBucket(String key) {
        return buckets.computeIfAbsent(key, k -> {
            // 定义令牌桶的带宽，这里使用周期性补充策略，模拟滚动窗口
            // 例如：每1秒补充10个令牌，桶容量为10个令牌
            Bandwidth limit = Bandwidth.builder()
                    .capacity(capacity) // 令牌桶容量
                    .refillGreedy(refillTokens, refillInterval) // 每 refillInterval 补充 refillTokens 个令牌
                    .build();

            // 使用Redis作为后端存储
            return Bucket4jJedis.builder()
                    .withJedisPool(jedisPool)
                    .with :key(key.getBytes()) // 将限流键作为Redis的key
                    .build()
                    .addLimit(limit)
                    .build();
        });
    }

    /**
     * 尝试消耗令牌并返回限流结果
     * @param key 限流键
     * @param tokensToConsume 尝试消耗的令牌数
     * @return ConsumptionProbe 对象，包含是否成功和回退时间
     */
    public ConsumptionProbe tryConsume(String key, long tokensToConsume) {
        Bucket bucket = getOrCreateBucket(key);
        return bucket.tryConsumeAndReturnRemaining(tokensToConsume);
    }

    public void destroy() {
        if (jedisPool != null) {
            jedisPool.close();
        }
    }
}

public class RateLimiterDemo {

    public static void main(String[] args) throws InterruptedException {
        // 初始化限流器：桶容量10个令牌，每1秒补充10个令牌
        RedisRateLimiter limiter = new RedisRateLimiter("localhost", 6379, 10, 10, Duration.ofSeconds(1));

        String userId = "user:123";
        System.out.println("--- 第一次尝试，连续发送请求 ---");
        for (int i = 0; i < 15; i++) {
            ConsumptionProbe probe = limiter.tryConsume(userId, 1);
            if (probe.isConsumed()) {
                System.out.printf("请求 %d 成功! 剩余令牌: %d%n", i + 1, probe.getRemainingTokens());
            } else {
                long nanosToWait = probe.getNanosToWaitForRefill();
                long millisToWait = TimeUnit.NANOSECONDS.toMillis(nanosToWait);
                System.out.printf("请求 %d 被限流! 请等待 %d 毫秒后重试。%n", i + 1, millisToWait);
                // 实际应用中，这里可以抛出异常或返回特定状态码，并在响应头中包含 Retry-After
                // 例如：throw new TooManyRequestsException("Too many requests", millisToWait);
            }
            Thread.sleep(50); // 模拟请求间隔
        }

        System.out.println("\n--- 等待一段时间后，再次尝试 ---");
        Thread.sleep(2000); // 等待2秒，让令牌桶有时间补充

        for (int i = 0; i < 5; i++) {
            ConsumptionProbe probe = limiter.tryConsume(userId, 1);
            if (probe.isConsumed()) {
                System.out.printf("再次请求 %d 成功! 剩余令牌: %d%n", i + 1, probe.getRemainingTokens());
            } else {
                long nanosToWait = probe.getNanosToWaitForRefill();
                long millisToWait = TimeUnit.NANOSECONDS.toMillis(nanosToWait);
                System.out.printf("再次请求 %d 被限流! 请等待 %d 毫秒后重试。%n", i + 1, millisToWait);
            }
            Thread.sleep(50);
        }

        limiter.destroy();
    }
}