限制API接口访问速率

文章目录

免责声明：本人无意侵权，奈何找不到原文作者，也找不到网址，于是自己记录一下，如果有侵权之嫌，请联系我删除文章

依赖

<!-- https://mvnrepository.com/artifact/com.google.guava/guava --><dependency><groupId>com.google.guava</groupId><artifactId>guava</artifactId><version>32.1.3-jre</version></dependency>

注解

importjava.lang.annotation.ElementType;importjava.lang.annotation.Retention;importjava.lang.annotation.RetentionPolicy;importjava.lang.annotation.Target;@Target(value =ElementType.METHOD)@Retention(value =RetentionPolicy.RUNTIME)public@interfaceRateConfigAnno{StringlimitType();doublelimitCount()default5d;}

aop

importcn.hutool.core.thread.ThreadUtil;importcom.alibaba.fastjson2.JSONObject;importcom.google.common.util.concurrent.RateLimiter;importcom.tjbchtyw.tjflowcontrol.annocation.RateConfigAnno;importcom.tjbchtyw.tjflowcontrol.helper.RateLimitHelper;importjakarta.servlet.ServletOutputStream;importjakarta.servlet.http.HttpServletResponse;importorg.aspectj.lang.JoinPoint;importorg.aspectj.lang.annotation.Aspect;importorg.aspectj.lang.annotation.Before;importorg.slf4j.Logger;importorg.slf4j.LoggerFactory;importorg.springframework.stereotype.Component;importorg.springframework.web.context.request.RequestContextHolder;importorg.springframework.web.context.request.ServletRequestAttributes;importjava.io.IOException;importjava.lang.reflect.Method;importjava.nio.charset.StandardCharsets;importjava.util.Objects;@Aspect@ComponentpublicclassGuavaLimitAop{privatestaticfinalLogger logger =LoggerFactory.getLogger(GuavaLimitAop.class);@Before("@annotation(com.tjbchtyw.tjflowcontrol.annocation.RateConfigAnno)")publicvoidlimit(JoinPoint joinPoint){//1、获取当前的调用方法Method currentMethod =getCurrentMethod(joinPoint);if(Objects.isNull(currentMethod)){return;}//2、从方法注解定义上获取限流的类型String limitType = currentMethod.getAnnotation(RateConfigAnno.class).limitType();double limitCount = currentMethod.getAnnotation(RateConfigAnno.class).limitCount();//使用guava的令牌桶算法获取一个令牌，获取不到先等待RateLimiter rateLimiter =RateLimitHelper.getRateLimiter(limitType, limitCount);//  boolean b =true;boolean pass = rateLimiter.tryAcquire();if(pass){System.out.println("获取到令牌");}else{//重试 仅测试用 有优化方案可以放在评论区for(int i =0; i <5; i++){ThreadUtil.safeSleep(1000);System.out.println("第"+(i +1)+"次获取令牌");if(rateLimiter.tryAcquire())break;if(i ==4){System.out.println("在第"+(i +1)+"次后未获取到令牌 开始限流");HttpServletResponse resp =((ServletRequestAttributes)RequestContextHolder.getRequestAttributes()).getResponse();JSONObject jsonObject=newJSONObject();
                    jsonObject.put("success",false);
                    jsonObject.put("msg","限流中");try{output(resp, jsonObject.toJSONString());}catch(Exception e){
                        logger.error("error,e:{}",e);}}}}}privateMethodgetCurrentMethod(JoinPoint joinPoint){Method[] methods = joinPoint.getTarget().getClass().getMethods();Method target =null;for(Method method : methods){if(method.getName().equals(joinPoint.getSignature().getName())){
                target = method;break;}}return target;}publicvoidoutput(HttpServletResponse response,String msg)throwsIOException{
        response.setContentType("application/json;charset=UTF-8");ServletOutputStream outputStream =null;try{
            outputStream = response.getOutputStream();
            outputStream.write(msg.getBytes(StandardCharsets.UTF_8));}catch(IOException e){
            e.printStackTrace();}finally{
            outputStream.flush();
            outputStream.close();}}}

helper

importcom.google.common.util.concurrent.RateLimiter;importjava.util.HashMap;importjava.util.Map;publicclassRateLimitHelper{privateRateLimitHelper(){}privatestaticfinalMap<String,RateLimiter> rateMap =newHashMap<>();publicstaticRateLimitergetRateLimiter(String limitType,double limitCount ){RateLimiter rateLimiter = rateMap.get(limitType);if(rateLimiter ==null){
            rateLimiter =RateLimiter.create(limitCount);
            rateMap.put(limitType,rateLimiter);}return rateLimiter;}}

Test

@RestController@Tag(name ="测试Controller", description ="这是描述")@RequestMapping("/pred-api")publicclassFlowController{@AutowiredFlowContext flowContext;@PostMapping("/pdf/test")//Count 限制次数@RateConfigAnno(limitType ="makePdf",limitCount =15)@Operation(summary ="限流接口")publicStringflowCont(@Parameter(name ="pdfParam", description ="参数对象,type标识执行不同的策略")@RequestBodyParamParam){return"test";}}