This commit is contained in:
dushitaoyuan
2020-12-23 17:51:23 +08:00
2 changed files with 31 additions and 42 deletions

View File

@@ -61,18 +61,14 @@ public class RateLimitAspect {
RateLimitType type = rateLimit.type(); RateLimitType type = rateLimit.type();
String key = rateLimit.key(); String key = rateLimit.key();
if (type == null) { if (type == null) {
key = methodName; type = RateLimitType.METHOD;
if (LOG.isDebugEnabled()) {
LOG.debug("{}限流策略未定义,采用[{}]限流策略", methodName, RateLimitType.METHOD);
} }
} else {
switch (type) { switch (type) {
case IP: case IP:
String serviceKey = rateLimit.key(); if (key == null || key.isEmpty()) {
if (serviceKey == null || key.isEmpty()) {
key = RequestUtil.getRemoteIp() + "_" + methodName; key = RequestUtil.getRemoteIp() + "_" + methodName;
} else { } else {
key = RequestUtil.getRemoteIp() + "_" + serviceKey; key = RequestUtil.getRemoteIp() + "_" + key;
} }
break; break;
case METHOD: case METHOD:
@@ -84,7 +80,6 @@ public class RateLimitAspect {
key = "global"; key = "global";
break; break;
} }
}
if (LOG.isDebugEnabled()) { if (LOG.isDebugEnabled()) {
LOG.debug("采用[{}]限流策略,限流key:{}", type, key); LOG.debug("采用[{}]限流策略,限流key:{}", type, key);
} }

View File

@@ -31,12 +31,10 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
rateHolder.clear(); rateHolder.clear();
} }
RateLimiter rateLimiter = null; RateLimiter rateLimiter = null;
if (rateHolder.containsKey(key)) { if (!rateHolder.containsKey(key)) {
rateLimiter = rateHolder.get(key); rateHolder.putIfAbsent(key, RateLimiter.create(limit));
} else {
rateLimiter = RateLimiter.create(limit);
} }
rateHolder.putIfAbsent(key, rateLimiter); rateLimiter = rateHolder.get(key);
return rateLimiter.tryAcquire(permits); return rateLimiter.tryAcquire(permits);
} }
@@ -54,7 +52,9 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
countHolder.clear(); countHolder.clear();
} }
LongAdder longAdder = null; LongAdder longAdder = null;
if (countHolder.containsKey(key)) { if (!countHolder.containsKey(key)) {
countHolder.putIfAbsent(key, new LongAdder());
}
longAdder = countHolder.get(key); longAdder = countHolder.get(key);
if (longAdder.longValue() >= totalCount) { if (longAdder.longValue() >= totalCount) {
TOTAL_LIMIT_ZERO_FLAG.put(key); TOTAL_LIMIT_ZERO_FLAG.put(key);
@@ -65,11 +65,5 @@ public class GuavaRateLimiter extends AbstractRateLimiter {
return true; return true;
} }
longAdder = new LongAdder();
countHolder.putIfAbsent(key, longAdder);
countHolder.get(key).add(count);
return true;
}
} }