forked from nakatsusizuru/Pixiv-Illustration-Collection
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
限流模块
- Loading branch information
Showing
17 changed files
with
194 additions
and
142 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
98 changes: 56 additions & 42 deletions
98
src/main/java/dev/cheerfun/pixivic/ratelimit/aop/RateLimitProcessor.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,70 +1,84 @@ | ||
| package dev.cheerfun.pixivic.ratelimit.aop; | ||
|
|
||
| import com.google.common.collect.Maps; | ||
| import com.google.common.util.concurrent.RateLimiter; | ||
| import dev.cheerfun.pixivic.ratelimit.annotation.RateLimit; | ||
| import dev.cheerfun.pixivic.auth.constant.PermissionLevel; | ||
| import dev.cheerfun.pixivic.common.context.AppContext; | ||
| import dev.cheerfun.pixivic.ratelimit.exception.RateLimitException; | ||
| import io.github.bucket4j.*; | ||
| import lombok.RequiredArgsConstructor; | ||
| import lombok.extern.slf4j.Slf4j; | ||
| import org.aspectj.lang.ProceedingJoinPoint; | ||
| import org.aspectj.lang.Signature; | ||
| import org.aspectj.lang.annotation.Around; | ||
| import org.aspectj.lang.annotation.Aspect; | ||
| import org.aspectj.lang.annotation.Pointcut; | ||
| import org.aspectj.lang.reflect.MethodSignature; | ||
| import org.springframework.beans.factory.annotation.Autowired; | ||
| import org.springframework.core.annotation.Order; | ||
| import org.springframework.http.HttpStatus; | ||
| import org.springframework.http.ResponseEntity; | ||
| import org.springframework.stereotype.Component; | ||
| import org.springframework.web.servlet.HandlerInterceptor; | ||
|
|
||
| import java.time.Duration; | ||
| import java.util.Map; | ||
|
|
||
| import java.util.concurrent.ConcurrentHashMap; | ||
|
|
||
| /** | ||
| * @author echo huang | ||
| * @version 1.0 | ||
| * @date 2019-07-13 15:12 | ||
| * @description 限流监听器 | ||
| * @author OysterQAQ | ||
| * @version 2.0 | ||
| * @date 2019-11-22 15:12 | ||
| * @description 限流处理器 | ||
| */ | ||
| @Aspect | ||
| @Component | ||
| @Slf4j | ||
| public class RateLimitProcessor { | ||
|
|
||
| /** | ||
| * 用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter) | ||
| */ | ||
| private Map<String, RateLimiter> map = Maps.newConcurrentMap(); | ||
| private RateLimiter rateLimiter; | ||
| @RequiredArgsConstructor(onConstructor = @__(@Autowired)) | ||
| @Order(1) | ||
| public class RateLimitProcessor implements HandlerInterceptor { | ||
| private static final String USER_ID = "userId"; | ||
| private final static String PERMISSION_LEVEL = "permissionLevel"; | ||
| private final Map<String, Bucket> buckets = new ConcurrentHashMap<>(); | ||
| //未登录用户 | ||
| private final Bucket freeBucket = Bucket4j.builder() | ||
| .addLimit(Bandwidth.classic(10, Refill.intervally(10, Duration.ofMinutes(1)))) | ||
| .build(); | ||
|
|
||
| @Pointcut(value = "@annotation(dev.cheerfun.pixivic.ratelimit.annotation.RateLimit)") | ||
| @Pointcut(value = "@annotation(dev.cheerfun.pixivic.ratelimit.annotation.RateLimit)||@within(dev.cheerfun.pixivic.ratelimit.annotation.RateLimit)") | ||
| public void pointCut() { | ||
| } | ||
|
|
||
| @Around(value = "pointCut()") | ||
| public Object handleRateLimiter(ProceedingJoinPoint joinPoint) { | ||
| Object obj = null; | ||
| Signature signature = joinPoint.getSignature(); | ||
| MethodSignature methodSignature = (MethodSignature) signature; | ||
| RateLimit annotation = methodSignature.getMethod().getAnnotation(RateLimit.class); | ||
| double limitNum = annotation.limitNum(); | ||
| String functionName = methodSignature.getName(); | ||
| if (map.containsKey(functionName)) { | ||
| rateLimiter = map.get(functionName); | ||
| public ResponseEntity handleRateLimit(ProceedingJoinPoint joinPoint) throws Throwable { | ||
| Bucket requestBucket; | ||
| if (AppContext.get() != null) { | ||
| Integer userId = (Integer) AppContext.get().get(USER_ID); | ||
| Integer permissionLevel = (Integer) AppContext.get().get(PERMISSION_LEVEL); | ||
| if (permissionLevel == PermissionLevel.VIP) { | ||
| requestBucket = this.buckets.computeIfAbsent(userId.toString(), key -> premiumBucket()); | ||
| } else { | ||
| requestBucket = this.buckets.computeIfAbsent(userId.toString(), key -> standardBucket()); | ||
| } | ||
| } else { | ||
| //每秒允许多少请求limitNum | ||
| map.put(functionName, RateLimiter.create(limitNum)); | ||
| rateLimiter = map.get(functionName); | ||
| requestBucket = this.freeBucket; | ||
| } | ||
|
|
||
| if (rateLimiter.tryAcquire()) { | ||
| //执行方法 | ||
| try { | ||
| obj = joinPoint.proceed(); | ||
| } catch (Throwable throwable) { | ||
| throwable.printStackTrace(); | ||
| } | ||
| } else { | ||
| //拒绝了请求（服务降级） | ||
| //TODO 拒绝后提示 | ||
| //throw new VisitOftenException(HttpStatus.FORBIDDEN.value(), HttpStatus.FORBIDDEN.getReasonPhrase()); | ||
| /* if (userId != null && permissionLevel != null) { | ||
| }*/ | ||
| ConsumptionProbe probe = requestBucket.tryConsumeAndReturnRemaining(1); | ||
| if (probe.isConsumed()) { | ||
| return (ResponseEntity) joinPoint.proceed(); | ||
| } | ||
| throw new RateLimitException(HttpStatus.TOO_MANY_REQUESTS, "请求过于频繁"); | ||
| } | ||
|
|
||
| private static Bucket standardBucket() { | ||
| return Bucket4j.builder() | ||
| .addLimit(Bandwidth.classic(50, Refill.intervally(50, Duration.ofMinutes(1)))) | ||
| .build(); | ||
| } | ||
|
|
||
| return obj; | ||
| private static Bucket premiumBucket() { | ||
| return Bucket4j.builder() | ||
| .addLimit(Bandwidth.classic(100, Refill.intervally(100, Duration.ofMinutes(1)))) | ||
| .build(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.