限制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>

注解

import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;@Target(value = ElementType.METHOD)
@Retention(value = RetentionPolicy.RUNTIME)
public @interface RateConfigAnno {String limitType();double limitCount() default 5d;
}

aop

import cn.hutool.core.thread.ThreadUtil;
import com.alibaba.fastjson2.JSONObject;
import com.google.common.util.concurrent.RateLimiter;
import com.tjbchtyw.tjflowcontrol.annocation.RateConfigAnno;
import com.tjbchtyw.tjflowcontrol.helper.RateLimitHelper;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.http.HttpServletResponse;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;import java.io.IOException;
import java.lang.reflect.Method;
import java.nio.charset.StandardCharsets;
import java.util.Objects;@Aspect
@Component
public class GuavaLimitAop {private static final Logger logger = LoggerFactory.getLogger(GuavaLimitAop.class);@Before("@annotation(com.tjbchtyw.tjflowcontrol.annocation.RateConfigAnno)")public void limit(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=new JSONObject();jsonObject.put("success",false);jsonObject.put("msg","限流中");try {output(resp, jsonObject.toJSONString());}catch (Exception e){logger.error("error,e:{}",e);}}}}}private Method getCurrentMethod(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;}public void output(HttpServletResponse response, String msg) throws IOException {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

import com.google.common.util.concurrent.RateLimiter;import java.util.HashMap;
import java.util.Map;public class RateLimitHelper {private RateLimitHelper(){}private static final Map<String,RateLimiter> rateMap = new HashMap<>();public static RateLimiter getRateLimiter(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")public class FlowController {@AutowiredFlowContext flowContext;@PostMapping("/pdf/test")//Count 限制次数@RateConfigAnno(limitType = "makePdf",limitCount = 15)@Operation(summary = "限流接口")public String flowCont(@Parameter(name = "pdfParam", description = "参数对象,type标识执行不同的策略") @RequestBody Param Param) {return "test";}}