From dcaeb18870ecdf319dbdb03820405a745d4eeb32 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=96=AF=E7=8B=82=E7=9A=84=E7=8B=AE=E5=AD=90li?=
 <15040126243@163.com>
Date: Fri, 17 Sep 2021 17:22:45 +0800
Subject: [PATCH] =?UTF-8?q?add=20=E5=A2=9E=E5=8A=A0=20=E9=9B=86=E7=BE=A4?=
 =?UTF-8?q?=E9=99=90=E6=B5=81=E5=8A=9F=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../com/ruoyi/common/enums/LimitType.java     |  7 ++++-
 .../framework/aspectj/RateLimiterAspect.java  | 28 +++++++++++++++----
 2 files changed, 28 insertions(+), 7 deletions(-)

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