diff --git a/yudao-framework/yudao-spring-boot-starter-protection/pom.xml b/yudao-framework/yudao-spring-boot-starter-protection/pom.xml index a93991ff5..bbb5b12eb 100644 --- a/yudao-framework/yudao-spring-boot-starter-protection/pom.xml +++ b/yudao-framework/yudao-spring-boot-starter-protection/pom.xml @@ -20,7 +20,7 @@ cn.iocoder.boot yudao-spring-boot-starter-web - provided + provided diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/idempotent/core/aop/IdempotentAspect.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/idempotent/core/aop/IdempotentAspect.java index 2d8c76d55..11ff76576 100644 --- a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/idempotent/core/aop/IdempotentAspect.java +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/idempotent/core/aop/IdempotentAspect.java @@ -37,7 +37,7 @@ public class IdempotentAspect { } @Around(value = "@annotation(idempotent)") - public Object beforePointCut(ProceedingJoinPoint joinPoint, Idempotent idempotent) throws Throwable { + public Object aroundPointCut(ProceedingJoinPoint joinPoint, Idempotent idempotent) throws Throwable { // 获得 IdempotentKeyResolver IdempotentKeyResolver keyResolver = keyResolvers.get(idempotent.keyResolver()); Assert.notNull(keyResolver, "找不到对应的 IdempotentKeyResolver"); @@ -48,7 +48,7 @@ public class IdempotentAspect { boolean success = idempotentRedisDAO.setIfAbsent(key, idempotent.timeout(), idempotent.timeUnit()); // 锁定失败，抛出异常 if (!success) { - log.info("[beforePointCut][方法({}) 参数({}) 存在重复请求]", joinPoint.getSignature().toString(), joinPoint.getArgs()); + log.info("[aroundPointCut][方法({}) 参数({}) 存在重复请求]", joinPoint.getSignature().toString(), joinPoint.getArgs()); throw new ServiceException(GlobalErrorCodeConstants.REPEATED_REQUESTS.getCode(), idempotent.message()); } diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/config/YudaoRateLimiterConfiguration.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/config/YudaoRateLimiterConfiguration.java new file mode 100644 index 000000000..68b910fea --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/config/YudaoRateLimiterConfiguration.java @@ -0,0 +1,55 @@ +package cn.iocoder.yudao.framework.ratelimiter.config; + +import cn.iocoder.yudao.framework.ratelimiter.core.aop.RateLimiterAspect; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl.*; +import cn.iocoder.yudao.framework.ratelimiter.core.redis.RateLimiterRedisDAO; +import cn.iocoder.yudao.framework.redis.config.YudaoRedisAutoConfiguration; +import org.redisson.api.RedissonClient; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.context.annotation.Bean; + +import java.util.List; + +@AutoConfiguration(after = YudaoRedisAutoConfiguration.class) +public class YudaoRateLimiterConfiguration { + + @Bean + public RateLimiterAspect rateLimiterAspect(List keyResolvers, RateLimiterRedisDAO rateLimiterRedisDAO) { + return new RateLimiterAspect(keyResolvers, rateLimiterRedisDAO); + } + + @Bean + @SuppressWarnings("SpringJavaInjectionPointsAutowiringInspection") + public RateLimiterRedisDAO rateLimiterRedisDAO(RedissonClient redissonClient) { + return new RateLimiterRedisDAO(redissonClient); + } + + // ========== 各种 RateLimiterRedisDAO Bean ========== + + @Bean + public DefaultRateLimiterKeyResolver defaultRateLimiterKeyResolver() { + return new DefaultRateLimiterKeyResolver(); + } + + @Bean + public UserRateLimiterKeyResolver userRateLimiterKeyResolver() { + return new UserRateLimiterKeyResolver(); + } + + @Bean + public ClientIpRateLimiterKeyResolver clientIpRateLimiterKeyResolver() { + return new ClientIpRateLimiterKeyResolver(); + } + + @Bean + public ServerNodeRateLimiterKeyResolver serverNodeRateLimiterKeyResolver() { + return new ServerNodeRateLimiterKeyResolver(); + } + + @Bean + public ExpressionRateLimiterKeyResolver expressionRateLimiterKeyResolver() { + return new ExpressionRateLimiterKeyResolver(); + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/annotation/RateLimiter.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/annotation/RateLimiter.java new file mode 100644 index 000000000..417c4d642 --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/annotation/RateLimiter.java @@ -0,0 +1,62 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.annotation; + +import cn.iocoder.yudao.framework.common.exception.enums.GlobalErrorCodeConstants; +import cn.iocoder.yudao.framework.idempotent.core.keyresolver.impl.ExpressionIdempotentKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl.ClientIpRateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl.DefaultRateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl.ServerNodeRateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl.UserRateLimiterKeyResolver; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; +import java.util.concurrent.TimeUnit; + +/** + * 限流注解 + * + * @author 芋道源码 + */ +@Target({ElementType.METHOD}) +@Retention(RetentionPolicy.RUNTIME) +public @interface RateLimiter { + + /** + * 限流的时间，默认为 1 秒 + */ + int time() default 1; + /** + * 时间单位，默认为 SECONDS 秒 + */ + TimeUnit timeUnit() default TimeUnit.SECONDS; + + /** + * 限流次数 + */ + int count() default 100; + + /** + * 提示信息，请求过快的提示 + * + * @see GlobalErrorCodeConstants#TOO_MANY_REQUESTS + */ + String message() default ""; // 为空时，使用 TOO_MANY_REQUESTS 错误提示 + + /** + * 使用的 Key 解析器 + * + * @see DefaultRateLimiterKeyResolver 全局级别 + * @see UserRateLimiterKeyResolver 用户 ID 级别 + * @see ClientIpRateLimiterKeyResolver 用户 IP 级别 + * @see ServerNodeRateLimiterKeyResolver 服务器 Node 级别 + * @see ExpressionIdempotentKeyResolver 自定义表达式，通过 {@link #keyArg()} 计算 + */ + Class keyResolver() default DefaultRateLimiterKeyResolver.class; + /** + * 使用的 Key 参数 + */ + String keyArg() default ""; + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/aop/RateLimiterAspect.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/aop/RateLimiterAspect.java new file mode 100644 index 000000000..6ede62bea --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/aop/RateLimiterAspect.java @@ -0,0 +1,60 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.aop; + +import cn.hutool.core.util.StrUtil; +import cn.iocoder.yudao.framework.common.exception.ServiceException; +import cn.iocoder.yudao.framework.common.exception.enums.GlobalErrorCodeConstants; +import cn.iocoder.yudao.framework.common.util.collection.CollectionUtils; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import cn.iocoder.yudao.framework.ratelimiter.core.redis.RateLimiterRedisDAO; +import lombok.extern.slf4j.Slf4j; +import org.aspectj.lang.JoinPoint; +import org.aspectj.lang.annotation.Aspect; +import org.aspectj.lang.annotation.Before; +import org.springframework.util.Assert; + +import java.util.List; +import java.util.Map; + +/** + * 拦截声明了 {@link RateLimiter} 注解的方法，实现限流操作 + * + * @author 芋道源码 + */ +@Aspect +@Slf4j +public class RateLimiterAspect { + + /** + * RateLimiterKeyResolver 集合 + */ + private final Map, RateLimiterKeyResolver> keyResolvers; + + private final RateLimiterRedisDAO rateLimiterRedisDAO; + + public RateLimiterAspect(List keyResolvers, RateLimiterRedisDAO rateLimiterRedisDAO) { + this.keyResolvers = CollectionUtils.convertMap(keyResolvers, RateLimiterKeyResolver::getClass); + this.rateLimiterRedisDAO = rateLimiterRedisDAO; + } + + @Before("@annotation(rateLimiter)") + public void beforePointCut(JoinPoint joinPoint, RateLimiter rateLimiter) { + // 获得 IdempotentKeyResolver 对象 + RateLimiterKeyResolver keyResolver = keyResolvers.get(rateLimiter.keyResolver()); + Assert.notNull(keyResolver, "找不到对应的 RateLimiterKeyResolver"); + // 解析 Key + String key = keyResolver.resolver(joinPoint, rateLimiter); + + // 获取 1 次限流 + boolean success = rateLimiterRedisDAO.tryAcquire(key, + rateLimiter.count(), rateLimiter.time(), rateLimiter.timeUnit()); + if (!success) { + log.info("[beforePointCut][方法({}) 参数({}) 请求过于频繁]", joinPoint.getSignature().toString(), joinPoint.getArgs()); + String message = StrUtil.blankToDefault(rateLimiter.message(), + GlobalErrorCodeConstants.TOO_MANY_REQUESTS.getMsg()); + throw new ServiceException(GlobalErrorCodeConstants.TOO_MANY_REQUESTS.getCode(), message); + } + } + +} + diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/RateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/RateLimiterKeyResolver.java new file mode 100644 index 000000000..44d7bdff5 --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/RateLimiterKeyResolver.java @@ -0,0 +1,22 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver; + +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import org.aspectj.lang.JoinPoint; + +/** + * 限流 Key 解析器接口 + * + * @author 芋道源码 + */ +public interface RateLimiterKeyResolver { + + /** + * 解析一个 Key + * + * @param rateLimiter 限流注解 + * @param joinPoint AOP 切面 + * @return Key + */ + String resolver(JoinPoint joinPoint, RateLimiter rateLimiter); + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ClientIpRateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ClientIpRateLimiterKeyResolver.java new file mode 100644 index 000000000..8d6253caa --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ClientIpRateLimiterKeyResolver.java @@ -0,0 +1,27 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl; + +import cn.hutool.core.util.StrUtil; +import cn.hutool.crypto.SecureUtil; +import cn.iocoder.yudao.framework.common.util.servlet.ServletUtils; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import org.aspectj.lang.JoinPoint; + +/** + * IP 级别的限流 Key 解析器，使用方法名 + 方法参数 + IP，组装成一个 Key + * + * 为了避免 Key 过长，使用 MD5 进行“压缩” + * + * @author 芋道源码 + */ +public class ClientIpRateLimiterKeyResolver implements RateLimiterKeyResolver { + + @Override + public String resolver(JoinPoint joinPoint, RateLimiter rateLimiter) { + String methodName = joinPoint.getSignature().toString(); + String argsStr = StrUtil.join(",", joinPoint.getArgs()); + String clientIp = ServletUtils.getClientIP(); + return SecureUtil.md5(methodName + argsStr + clientIp); + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/DefaultRateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/DefaultRateLimiterKeyResolver.java new file mode 100644 index 000000000..236ea45cb --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/DefaultRateLimiterKeyResolver.java @@ -0,0 +1,25 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl; + +import cn.hutool.core.util.StrUtil; +import cn.hutool.crypto.SecureUtil; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import org.aspectj.lang.JoinPoint; + +/** + * 默认（全局级别）限流 Key 解析器，使用方法名 + 方法参数，组装成一个 Key + * + * 为了避免 Key 过长，使用 MD5 进行“压缩” + * + * @author 芋道源码 + */ +public class DefaultRateLimiterKeyResolver implements RateLimiterKeyResolver { + + @Override + public String resolver(JoinPoint joinPoint, RateLimiter rateLimiter) { + String methodName = joinPoint.getSignature().toString(); + String argsStr = StrUtil.join(",", joinPoint.getArgs()); + return SecureUtil.md5(methodName + argsStr); + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ExpressionRateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ExpressionRateLimiterKeyResolver.java new file mode 100644 index 000000000..118581e35 --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ExpressionRateLimiterKeyResolver.java @@ -0,0 +1,64 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl; + +import cn.hutool.core.util.ArrayUtil; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import org.aspectj.lang.JoinPoint; +import org.aspectj.lang.reflect.MethodSignature; +import org.springframework.core.DefaultParameterNameDiscoverer; +import org.springframework.core.ParameterNameDiscoverer; +import org.springframework.expression.Expression; +import org.springframework.expression.ExpressionParser; +import org.springframework.expression.spel.standard.SpelExpressionParser; +import org.springframework.expression.spel.support.StandardEvaluationContext; + +import java.lang.reflect.Method; + +/** + * 基于 Spring EL 表达式的 {@link RateLimiterKeyResolver} 实现类 + * + * @author 芋道源码 + */ +public class ExpressionRateLimiterKeyResolver implements RateLimiterKeyResolver { + + private final ParameterNameDiscoverer parameterNameDiscoverer = new DefaultParameterNameDiscoverer(); + + private final ExpressionParser expressionParser = new SpelExpressionParser(); + + @Override + public String resolver(JoinPoint joinPoint, RateLimiter rateLimiter) { + // 获得被拦截方法参数名列表 + Method method = getMethod(joinPoint); + Object[] args = joinPoint.getArgs(); + String[] parameterNames = this.parameterNameDiscoverer.getParameterNames(method); + // 准备 Spring EL 表达式解析的上下文 + StandardEvaluationContext evaluationContext = new StandardEvaluationContext(); + if (ArrayUtil.isNotEmpty(parameterNames)) { + for (int i = 0; i < parameterNames.length; i++) { + evaluationContext.setVariable(parameterNames[i], args[i]); + } + } + + // 解析参数 + Expression expression = expressionParser.parseExpression(rateLimiter.keyArg()); + return expression.getValue(evaluationContext, String.class); + } + + private static Method getMethod(JoinPoint point) { + // 处理，声明在类上的情况 + MethodSignature signature = (MethodSignature) point.getSignature(); + Method method = signature.getMethod(); + if (!method.getDeclaringClass().isInterface()) { + return method; + } + + // 处理，声明在接口上的情况 + try { + return point.getTarget().getClass().getDeclaredMethod( + point.getSignature().getName(), method.getParameterTypes()); + } catch (NoSuchMethodException e) { + throw new RuntimeException(e); + } + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ServerNodeRateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ServerNodeRateLimiterKeyResolver.java new file mode 100644 index 000000000..300a4d2f1 --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/ServerNodeRateLimiterKeyResolver.java @@ -0,0 +1,27 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl; + +import cn.hutool.core.util.StrUtil; +import cn.hutool.crypto.SecureUtil; +import cn.hutool.system.SystemUtil; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import org.aspectj.lang.JoinPoint; + +/** + * Server 节点级别的限流 Key 解析器，使用方法名 + 方法参数 + IP，组装成一个 Key + * + * 为了避免 Key 过长，使用 MD5 进行“压缩” + * + * @author 芋道源码 + */ +public class ServerNodeRateLimiterKeyResolver implements RateLimiterKeyResolver { + + @Override + public String resolver(JoinPoint joinPoint, RateLimiter rateLimiter) { + String methodName = joinPoint.getSignature().toString(); + String argsStr = StrUtil.join(",", joinPoint.getArgs()); + String serverNode = String.format("%s@%d", SystemUtil.getHostInfo().getAddress(), SystemUtil.getCurrentPID()); + return SecureUtil.md5(methodName + argsStr + serverNode); + } + +} \ No newline at end of file diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/UserRateLimiterKeyResolver.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/UserRateLimiterKeyResolver.java new file mode 100644 index 000000000..a8d1c3a98 --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/keyresolver/impl/UserRateLimiterKeyResolver.java @@ -0,0 +1,28 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.impl; + +import cn.hutool.core.util.StrUtil; +import cn.hutool.crypto.SecureUtil; +import cn.iocoder.yudao.framework.ratelimiter.core.annotation.RateLimiter; +import cn.iocoder.yudao.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; +import cn.iocoder.yudao.framework.web.core.util.WebFrameworkUtils; +import org.aspectj.lang.JoinPoint; + +/** + * 用户级别的限流 Key 解析器，使用方法名 + 方法参数 + userId + userType，组装成一个 Key + * + * 为了避免 Key 过长，使用 MD5 进行“压缩” + * + * @author 芋道源码 + */ +public class UserRateLimiterKeyResolver implements RateLimiterKeyResolver { + + @Override + public String resolver(JoinPoint joinPoint, RateLimiter rateLimiter) { + String methodName = joinPoint.getSignature().toString(); + String argsStr = StrUtil.join(",", joinPoint.getArgs()); + Long userId = WebFrameworkUtils.getLoginUserId(); + Integer userType = WebFrameworkUtils.getLoginUserType(); + return SecureUtil.md5(methodName + argsStr + userId + userType); + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/redis/RateLimiterRedisDAO.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/redis/RateLimiterRedisDAO.java new file mode 100644 index 000000000..fc1378f3b --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/core/redis/RateLimiterRedisDAO.java @@ -0,0 +1,60 @@ +package cn.iocoder.yudao.framework.ratelimiter.core.redis; + +import lombok.AllArgsConstructor; +import org.redisson.api.*; + +import java.util.Objects; +import java.util.concurrent.TimeUnit; + +/** + * 限流 Redis DAO + * + * @author 芋道源码 + */ +@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; + } + +} diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/package-info.java b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/package-info.java new file mode 100644 index 000000000..36a408c8e --- /dev/null +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/java/cn/iocoder/yudao/framework/ratelimiter/package-info.java @@ -0,0 +1,4 @@ +/** + * 限流组件，基于 Redisson {@link org.redisson.api.RRateLimiter} 限流实现 + */ +package cn.iocoder.yudao.framework.ratelimiter; \ No newline at end of file diff --git a/yudao-framework/yudao-spring-boot-starter-protection/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/yudao-framework/yudao-spring-boot-starter-protection/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index be5c0990d..d7cd3a883 100644 --- a/yudao-framework/yudao-spring-boot-starter-protection/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/yudao-framework/yudao-spring-boot-starter-protection/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -1,2 +1,3 @@ cn.iocoder.yudao.framework.idempotent.config.YudaoIdempotentConfiguration -cn.iocoder.yudao.framework.lock4j.config.YudaoLock4jConfiguration \ No newline at end of file +cn.iocoder.yudao.framework.lock4j.config.YudaoLock4jConfiguration +cn.iocoder.yudao.framework.ratelimiter.config.YudaoRateLimiterConfiguration \ No newline at end of file