package com.iailab.framework.ratelimiter.core.redis; import lombok.AllArgsConstructor; import org.redisson.api.*; import java.util.Objects; import java.util.concurrent.TimeUnit; /** * 限流 Redis DAO * * @author iailab */ @AllArgsConstructor public class RateLimiterRedisDAO { /** * 限流操作 * * KEY 格式:rate_limiter:%s // 参数为 uuid * VALUE 格式:String * 过期时间:不固定 */ private static final String RATE_LIMITER = "rate_limiter:%s"; private final RedissonClient redissonClient; public Boolean tryAcquire(String key, int count, int time, TimeUnit timeUnit) { // 1. 获得 RRateLimiter,并设置 rate 速率 RRateLimiter rateLimiter = getRRateLimiter(key, count, time, timeUnit); // 2. 尝试获取 1 个 return rateLimiter.tryAcquire(); } private static String formatKey(String key) { return String.format(RATE_LIMITER, key); } private RRateLimiter getRRateLimiter(String key, long count, int time, TimeUnit timeUnit) { String redisKey = formatKey(key); RRateLimiter rateLimiter = redissonClient.getRateLimiter(redisKey); long rateInterval = timeUnit.toSeconds(time); // 1. 如果不存在,设置 rate 速率 RateLimiterConfig config = rateLimiter.getConfig(); if (config == null) { rateLimiter.trySetRate(RateType.OVERALL, count, rateInterval, RateIntervalUnit.SECONDS); return rateLimiter; } // 2. 如果存在,并且配置相同,则直接返回 if (config.getRateType() == RateType.OVERALL && Objects.equals(config.getRate(), count) && Objects.equals(config.getRateInterval(), TimeUnit.SECONDS.toMillis(rateInterval))) { return rateLimiter; } // 3. 如果存在,并且配置不同,则进行新建 rateLimiter.setRate(RateType.OVERALL, count, rateInterval, RateIntervalUnit.SECONDS); return rateLimiter; } }