diff --git a/ruoyi-common/src/main/java/com/ruoyi/common/enums/LimitType.java b/ruoyi-common/src/main/java/com/ruoyi/common/enums/LimitType.java index ad01659dc..02c2b3d30 100644 --- a/ruoyi-common/src/main/java/com/ruoyi/common/enums/LimitType.java +++ b/ruoyi-common/src/main/java/com/ruoyi/common/enums/LimitType.java @@ -16,5 +16,10 @@ public enum LimitType /** * 根据请求者IP进行限流 */ - IP + IP, + + /** + * 实例限流(集群多后端实例) + */ + CLUSTER } diff --git a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java index e63e47b1d..93502c04e 100644 --- a/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java +++ b/ruoyi-framework/src/main/java/com/ruoyi/framework/aspectj/RateLimiterAspect.java @@ -4,6 +4,7 @@ import com.ruoyi.common.annotation.RateLimiter; import com.ruoyi.common.enums.LimitType; import com.ruoyi.common.exception.ServiceException; import com.ruoyi.common.utils.RedisUtils; +import com.ruoyi.common.utils.ServletUtils; import lombok.extern.slf4j.Slf4j; import org.aspectj.lang.JoinPoint; import org.aspectj.lang.Signature; @@ -34,20 +35,19 @@ public class RateLimiterAspect { @Before("rateLimiterPointCut()") public void doBefore(JoinPoint point) throws Throwable { RateLimiter rateLimiter = getAnnotationRateLimiter(point); - String key = rateLimiter.key(); int time = rateLimiter.time(); int count = rateLimiter.count(); - + String combineKey = getCombineKey(rateLimiter, point); try { RateType rateType = RateType.OVERALL; - if (rateLimiter.limitType() == LimitType.IP) { + if (rateLimiter.limitType() == LimitType.CLUSTER) { rateType = RateType.PER_CLIENT; } - // 返回 false 说明 获取令牌失败 - if (!RedisUtils.rateLimiter(key, rateType, count, time)) { + long number = RedisUtils.rateLimiter(combineKey, rateType, count, time); + if (number == -1) { throw new ServiceException("访问过于频繁,请稍后再试"); } - log.info("限制请求'{}',缓存key'{}'", count, key); + log.info("限制令牌 => {}, 剩余令牌 => {}, 缓存key => '{}'", count, number, combineKey); } catch (ServiceException e) { throw e; } catch (Exception e) { @@ -69,4 +69,20 @@ public class RateLimiterAspect { return null; } + public String getCombineKey(RateLimiter rateLimiter, JoinPoint point) { + StringBuilder stringBuffer = new StringBuilder(rateLimiter.key()); + if (rateLimiter.limitType() == LimitType.IP) { + // 获取请求ip + stringBuffer.append(ServletUtils.getClientIP()).append("-"); + } else if (rateLimiter.limitType() == LimitType.CLUSTER){ + // 获取客户端实例id + stringBuffer.append(RedisUtils.getClientId()).append("-"); + } + MethodSignature signature = (MethodSignature) point.getSignature(); + Method method = signature.getMethod(); + Class targetClass = method.getDeclaringClass(); + stringBuffer.append(targetClass.getName()).append("-").append(method.getName()); + return stringBuffer.toString(); + } + }