From e305716c7550a4121114d7b8e8c6289425fa05b5 Mon Sep 17 00:00:00 2001
From: dushitaoyuan <dushitaoyuan@gmail.com>
Date: Sun, 20 Dec 2020 16:12:54 +0800
Subject: [PATCH] add

---
 .../security/RateLimitAspect.java             | 43 ++++++++-----------
 .../security/ratelimit/GuavaRateLimiter.java  | 30 ++++++-------
 2 files changed, 31 insertions(+), 42 deletions(-)

diff --git a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
index ff274a9..cdbb594 100644
--- a/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
+++ b/src/main/java/com/taoyuanx/securitydemo/security/RateLimitAspect.java
@@ -61,29 +61,24 @@ public class RateLimitAspect {
         RateLimitType type = rateLimit.type();
         String key = rateLimit.key();
         if (type == null) {
-            key = methodName;
-            if (LOG.isDebugEnabled()) {
-                LOG.debug("{}限流策略未定义,采用[{}]限流策略", methodName, RateLimitType.METHOD);
-            }
-        } else {
-            switch (type) {
-                case IP:
-                    String serviceKey = rateLimit.key();
-                    if (serviceKey == null || key.isEmpty()) {
-                        key = RequestUtil.getRemoteIp() + "_" + methodName;
-                    } else {
-                        key = RequestUtil.getRemoteIp() + "_" + serviceKey;
-                    }
-                    break;
-                case METHOD:
-                    key = methodName;
-                    break;
-                case SERVICE_KEY:
-                    break;
-                case GLOBAL:
-                    key = "global";
-                    break;
-            }
+            type = RateLimitType.METHOD;
+        }
+        switch (type) {
+            case IP:
+                if (key == null || key.isEmpty()) {
+                    key = RequestUtil.getRemoteIp() + "_" + methodName;
+                } else {
+                    key = RequestUtil.getRemoteIp() + "_" + key;
+                }
+                break;
+            case METHOD:
+                key = methodName;
+                break;
+            case SERVICE_KEY:
+                break;
+            case GLOBAL:
+                key = "global";
+                break;
         }
         if (LOG.isDebugEnabled()) {
             LOG.debug("采用[{}]限流策略,限流key:{}", type, key);
@@ -102,4 +97,4 @@ public class RateLimitAspect {
     }
 
 
-}
\ No newline at end of file
+}
diff --git a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java
index fb7df47..6ce5a0e 100644
--- a/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java
+++ b/src/main/java/com/taoyuanx/securitydemo/security/ratelimit/GuavaRateLimiter.java
@@ -31,12 +31,10 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
             rateHolder.clear();
         }
         RateLimiter rateLimiter = null;
-        if (rateHolder.containsKey(key)) {
-            rateLimiter = rateHolder.get(key);
-        } else {
-            rateLimiter = RateLimiter.create(limit);
+        if (!rateHolder.containsKey(key)) {
+            rateHolder.putIfAbsent(key, RateLimiter.create(limit));
         }
-        rateHolder.putIfAbsent(key, rateLimiter);
+        rateLimiter = rateHolder.get(key);
         return rateLimiter.tryAcquire(permits);
     }
 
@@ -54,20 +52,16 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
             countHolder.clear();
         }
         LongAdder longAdder = null;
-        if (countHolder.containsKey(key)) {
-            longAdder = countHolder.get(key);
-            if (longAdder.longValue() >=totalCount) {
-                TOTAL_LIMIT_ZERO_FLAG.put(key);
-                countHolder.remove(key);
-                return false;
-            }
-            longAdder.add(count);
-            return true;
+        if (!countHolder.containsKey(key)) {
+            countHolder.putIfAbsent(key, new LongAdder());
         }
-
-        longAdder = new LongAdder();
-        countHolder.putIfAbsent(key, longAdder);
-        countHolder.get(key).add(count);
+        longAdder = countHolder.get(key);
+        if (longAdder.longValue() >= totalCount) {
+            TOTAL_LIMIT_ZERO_FLAG.put(key);
+            countHolder.remove(key);
+            return false;
+        }
+        longAdder.add(count);
         return true;
     }