| 提交 | 用户 | 时间 | ||
| e7c126 | 1 | package com.iailab.framework.ratelimiter.core.redis; |
| H | 2 | |
| 3 | import lombok.AllArgsConstructor; | |
| 4 | import org.redisson.api.*; | |
| 5 | ||
| 6 | import java.util.Objects; | |
| 7 | import java.util.concurrent.TimeUnit; | |
| 8 | ||
| 9 | /** | |
| 10 | * 限流 Redis DAO | |
| 11 | * | |
| 12 | * @author iailab | |
| 13 | */ | |
| 14 | @AllArgsConstructor | |
| 15 | public class RateLimiterRedisDAO { | |
| 16 | ||
| 17 | /** | |
| 18 | * 限流操作 | |
| 19 | * | |
| 20 | * KEY 格式：rate_limiter:%s // 参数为 uuid | |
| 21 | * VALUE 格式：String | |
| 22 | * 过期时间：不固定 | |
| 23 | */ | |
| 24 | private static final String RATE_LIMITER = "rate_limiter:%s"; | |
| 25 | ||
| 26 | private final RedissonClient redissonClient; | |
| 27 | ||
| 28 | public Boolean tryAcquire(String key, int count, int time, TimeUnit timeUnit) { | |
| 29 | // 1. 获得 RRateLimiter，并设置 rate 速率 | |
| 30 | RRateLimiter rateLimiter = getRRateLimiter(key, count, time, timeUnit); | |
| 31 | // 2. 尝试获取 1 个 | |
| 32 | return rateLimiter.tryAcquire(); | |
| 33 | } | |
| 34 | ||
| 35 | private static String formatKey(String key) { | |
| 36 | return String.format(RATE_LIMITER, key); | |
| 37 | } | |
| 38 | ||
| 39 | private RRateLimiter getRRateLimiter(String key, long count, int time, TimeUnit timeUnit) { | |
| 40 | String redisKey = formatKey(key); | |
| 41 | RRateLimiter rateLimiter = redissonClient.getRateLimiter(redisKey); | |
| 42 | long rateInterval = timeUnit.toSeconds(time); | |
| 43 | // 1. 如果不存在，设置 rate 速率 | |
| 44 | RateLimiterConfig config = rateLimiter.getConfig(); | |
| 45 | if (config == null) { | |
| 46 | rateLimiter.trySetRate(RateType.OVERALL, count, rateInterval, RateIntervalUnit.SECONDS); | |
| 47 | return rateLimiter; | |
| 48 | } | |
| 49 | // 2. 如果存在，并且配置相同，则直接返回 | |
| 50 | if (config.getRateType() == RateType.OVERALL | |
| 51 | && Objects.equals(config.getRate(), count) | |
| 52 | && Objects.equals(config.getRateInterval(), TimeUnit.SECONDS.toMillis(rateInterval))) { | |
| 53 | return rateLimiter; | |
| 54 | } | |
| 55 | // 3. 如果存在，并且配置不同，则进行新建 | |
| 56 | rateLimiter.setRate(RateType.OVERALL, count, rateInterval, RateIntervalUnit.SECONDS); | |
| 57 | return rateLimiter; | |
| 58 | } | |
| 59 | ||
| 60 | } | |
