| 提交 | 用户 | 时间 | ||
| e7c126 | 1 | package com.iailab.framework.ratelimiter.core.aop; |
| H | 2 | |
| 3 | import cn.hutool.core.util.StrUtil; | |
| 4 | import com.iailab.framework.common.exception.ServiceException; | |
| 5 | import com.iailab.framework.common.exception.enums.GlobalErrorCodeConstants; | |
| 6 | import com.iailab.framework.common.util.collection.CollectionUtils; | |
| 7 | import com.iailab.framework.ratelimiter.core.annotation.RateLimiter; | |
| 8 | import com.iailab.framework.ratelimiter.core.keyresolver.RateLimiterKeyResolver; | |
| 9 | import com.iailab.framework.ratelimiter.core.redis.RateLimiterRedisDAO; | |
| 10 | import lombok.extern.slf4j.Slf4j; | |
| 11 | import org.aspectj.lang.JoinPoint; | |
| 12 | import org.aspectj.lang.annotation.Aspect; | |
| 13 | import org.aspectj.lang.annotation.Before; | |
| 14 | import org.springframework.util.Assert; | |
| 15 | ||
| 16 | import java.util.List; | |
| 17 | import java.util.Map; | |
| 18 | ||
| 19 | /** | |
| 20 | * 拦截声明了 {@link RateLimiter} 注解的方法，实现限流操作 | |
| 21 | * | |
| 22 | * @author iailab | |
| 23 | */ | |
| 24 | @Aspect | |
| 25 | @Slf4j | |
| 26 | public class RateLimiterAspect { | |
| 27 | ||
| 28 | /** | |
| 29 | * RateLimiterKeyResolver 集合 | |
| 30 | */ | |
| 31 | private final Map<Class<? extends RateLimiterKeyResolver>, RateLimiterKeyResolver> keyResolvers; | |
| 32 | ||
| 33 | private final RateLimiterRedisDAO rateLimiterRedisDAO; | |
| 34 | ||
| 35 | public RateLimiterAspect(List<RateLimiterKeyResolver> keyResolvers, RateLimiterRedisDAO rateLimiterRedisDAO) { | |
| 36 | this.keyResolvers = CollectionUtils.convertMap(keyResolvers, RateLimiterKeyResolver::getClass); | |
| 37 | this.rateLimiterRedisDAO = rateLimiterRedisDAO; | |
| 38 | } | |
| 39 | ||
| 40 | @Before("@annotation(rateLimiter)") | |
| 41 | public void beforePointCut(JoinPoint joinPoint, RateLimiter rateLimiter) { | |
| 42 | // 获得 IdempotentKeyResolver 对象 | |
| 43 | RateLimiterKeyResolver keyResolver = keyResolvers.get(rateLimiter.keyResolver()); | |
| 44 | Assert.notNull(keyResolver, "找不到对应的 RateLimiterKeyResolver"); | |
| 45 | // 解析 Key | |
| 46 | String key = keyResolver.resolver(joinPoint, rateLimiter); | |
| 47 | ||
| 48 | // 获取 1 次限流 | |
| 49 | boolean success = rateLimiterRedisDAO.tryAcquire(key, | |
| 50 | rateLimiter.count(), rateLimiter.time(), rateLimiter.timeUnit()); | |
| 51 | if (!success) { | |
| 52 | log.info("[beforePointCut][方法({}) 参数({}) 请求过于频繁]", joinPoint.getSignature().toString(), joinPoint.getArgs()); | |
| 53 | String message = StrUtil.blankToDefault(rateLimiter.message(), | |
| 54 | GlobalErrorCodeConstants.TOO_MANY_REQUESTS.getMsg()); | |
| 55 | throw new ServiceException(GlobalErrorCodeConstants.TOO_MANY_REQUESTS.getCode(), message); | |
| 56 | } | |
| 57 | } | |
| 58 | ||
| 59 | } | |
| 60 | ||
