SpringBoot自定义注解+AOP的方式实现限流

蔚落 2022-10-23 08:09 198阅读 0赞

1.添加依赖

>   
> <dependency>  
>     <groupId>org.aspectj</groupId>  
>     <artifactId>aspectjweaver</artifactId>  
> </dependency>
> 
> <dependency>  
>     <groupId>com.google.guava</groupId>  
>     <artifactId>guava</artifactId>  
>     <version>20.0</version>  
> </dependency>

2.添加注解RateLimitAspect

>     @Inherited
>     @Documented
>     @Target({ElementType.METHOD, ElementType.FIELD, ElementType.TYPE})
>     @Retention(RetentionPolicy.RUNTIME)
>     public @interface RateLimitAspect {
>         double limitNum() default 20;  //默认每秒放入桶中的token
>     }

3.自定义切面 RateLimitAop

>     import com.fasterxml.jackson.databind.ObjectMapper;
>     import com.google.common.util.concurrent.RateLimiter;
>     import groovy.util.logging.Slf4j;
>     import lombok.SneakyThrows;
>     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.context.annotation.Scope;
>     import org.springframework.stereotype.Component;
>     
>     import javax.servlet.ServletOutputStream;
>     import javax.servlet.http.HttpServletResponse;
>     import java.io.IOException;
>     import java.lang.reflect.Method;
>     import java.util.concurrent.ConcurrentHashMap;
>     
>     
>     @Component
>     @Scope
>     @Aspect
>     @Slf4j
>     public class RateLimitAop {
>         @Autowired
>         private HttpServletResponse response;
>     
>         private RateLimiter rateLimiter;
>         //用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
>         private ConcurrentHashMap<String, RateLimiter> map = new ConcurrentHashMap<>();
>     
>         private static ObjectMapper objectMapper = new ObjectMapper();
>     
>         //注解
>         @Pointcut("@annotation(com.dycx.dyyy.web.config.RateLimitAspect)")
>         public void serviceLimit() {
>     
>         }
>     
>         @SneakyThrows
>         @Around("serviceLimit()")
>         public Object around(ProceedingJoinPoint joinPoint) {
>             Object obj = null;
>             //获取拦截的方法名
>             Signature sig = joinPoint.getSignature();
>             //获取拦截的方法名
>             MethodSignature msig = (MethodSignature) sig;
>             //返回被织入增加处理目标对象
>             Object target = joinPoint.getTarget();
>             //为了获取注解信息
>             Method currentMethod = target.getClass().getMethod(msig.getName(), msig.getParameterTypes());
>             //获取注解信息
>             RateLimitAspect annotation = currentMethod.getAnnotation(RateLimitAspect.class);
>             //获取注解每秒加入桶中的token
>             double limitNum = annotation.limitNum();
>             // 注解所在方法名区分不同的限流策略
>             String functionName = msig.getName();
>             //获取rateLimiter
>             if (!map.containsKey(functionName)) {
>                 map.put(functionName, RateLimiter.create(limitNum));
>             }
>             rateLimiter = map.get(functionName);
>             try {
>                 if (rateLimiter.tryAcquire()) {
>                     //执行方法
>                     obj = joinPoint.proceed();
>                 } else {
>                     //拒绝了请求（服务降级）
>                     String result = objectMapper.writeValueAsString("系统繁忙");
>                     output(response, result);
>                 }
>             } catch (Throwable throwable) {
>                 throwable.printStackTrace();
>             }
>             return obj;
>     
>         }
>     
>         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("UTF-8"));
>             } catch (IOException e) {
>                 e.printStackTrace();
>             } finally {
>                 outputStream.flush();
>                 outputStream.close();
>             }
>         }
>     }
> 
>

4.添加在方法上面

>     @RestController
>     public class TestController {
>      
>         @RateLimitAspect
>         @RequestMapping("/test")
>         public String test(){
>             return "test"
>         }
>     }