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

1.添加依赖

 

    org.aspectj
    aspectjweaver

    com.google.guava
    guava
    20.0

 

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 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"
    }
}