Ver código fonte

add 增加 集群限流功能

疯狂的狮子li 3 anos atrás
pai
commit
dcaeb18870

+ 6 - 1
ruoyi-common/src/main/java/com/ruoyi/common/enums/LimitType.java

@@ -16,5 +16,10 @@ public enum LimitType
     /**
      * 根据请求者IP进行限流
      */
-    IP
+    IP,
+
+    /**
+     * 实例限流(集群多后端实例)
+     */
+    CLUSTER
 }

+ 22 - 6
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();
+    }
+
 }