|
|
@@ -11,6 +11,8 @@ import org.springframework.core.annotation.AnnotationUtils; |
|
|
|
import org.springframework.stereotype.Component; |
|
|
|
|
|
|
|
import java.lang.reflect.Method; |
|
|
|
import java.util.concurrent.ConcurrentHashMap; |
|
|
|
import java.util.concurrent.ConcurrentMap; |
|
|
|
|
|
|
|
/** |
|
|
|
* <p> |
|
|
@@ -24,7 +26,7 @@ import java.lang.reflect.Method; |
|
|
|
@Aspect |
|
|
|
@Component |
|
|
|
public class RateLimiterAspect { |
|
|
|
private static final com.google.common.util.concurrent.RateLimiter RATE_LIMITER = com.google.common.util.concurrent.RateLimiter.create(Double.MAX_VALUE); |
|
|
|
private static final ConcurrentMap<String, com.google.common.util.concurrent.RateLimiter> RATE_LIMITER_CACHE = new ConcurrentHashMap<>(); |
|
|
|
|
|
|
|
@Pointcut("@annotation(com.xkcoding.ratelimit.guava.annotation.RateLimiter)") |
|
|
|
public void rateLimit() { |
|
|
@@ -39,11 +41,14 @@ public class RateLimiterAspect { |
|
|
|
RateLimiter rateLimiter = AnnotationUtils.findAnnotation(method, RateLimiter.class); |
|
|
|
if (rateLimiter != null && rateLimiter.qps() > RateLimiter.NOT_LIMITED) { |
|
|
|
double qps = rateLimiter.qps(); |
|
|
|
log.debug("【{}】的QPS设置为: {}", method.getName(), qps); |
|
|
|
// 重新设置 QPS |
|
|
|
RATE_LIMITER.setRate(qps); |
|
|
|
if (RATE_LIMITER_CACHE.get(method.getName()) == null) { |
|
|
|
// 初始化 QPS |
|
|
|
RATE_LIMITER_CACHE.put(method.getName(), com.google.common.util.concurrent.RateLimiter.create(qps)); |
|
|
|
} |
|
|
|
|
|
|
|
log.debug("【{}】的QPS设置为: {}", method.getName(), RATE_LIMITER_CACHE.get(method.getName()).getRate()); |
|
|
|
// 尝试获取令牌 |
|
|
|
if (!RATE_LIMITER.tryAcquire(rateLimiter.timeout(), rateLimiter.timeUnit())) { |
|
|
|
if (RATE_LIMITER_CACHE.get(method.getName()) != null && !RATE_LIMITER_CACHE.get(method.getName()).tryAcquire(rateLimiter.timeout(), rateLimiter.timeUnit())) { |
|
|
|
throw new RuntimeException("手速太快了,慢点儿吧~"); |
|
|
|
} |
|
|
|
} |
|
|
|