使用 google guava 实现简单的限流

创建注解标记那些控制器需要限流

/**
 * 限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RequestRateLimiter {
    /**
     * 这里指吞吐率每秒多少许可数（通常是指QPS，每秒多少查询）
     * @return
     */
    double QPS() default 10D;

    /**
     * 获取令牌超时时间
     * @return
     */
    long acquireTokenTimeout() default 100;

    /**
     * 获取令牌超时时间单位：默认为 毫秒
     * @return
     */
    TimeUnit timeunit() default TimeUnit.MILLISECONDS;

    /**
     * 获取令牌失败提示
     * @return
     */
    String resMsg() default "请求过于繁忙请稍后再试！";
}

创建一个结果类

public class R<T> {

    private int code = 200;
    private String msg = "SUCCESS";
    private T data;

    public R(Builder<T> builder){
        this.code = builder.code;
        this.msg = builder.msg;
        this.data = builder.data;
    }

    public static class Builder<T>{
        private int code = 200;
        private String msg = "SUCCESS";
        private T data;

        public R ok(){
            return new R(this);
        }

        public R failed(){
            this.code = -1;
            this.msg = "failed";
            return new R(this);
        }

        public R consumerFailed(int code,String msg){
            this.code = code;
            this.msg = msg;
            return new R(this);
        }
        public Builder<T> setDate(T data){
            this.data = data;
            return this;
        }
    }

    public int getCode() {
        return code;
    }

    public void setCode(int code) {
        this.code = code;
    }

    public String getMsg() {
        return msg;
    }

    public void setMsg(String msg) {
        this.msg = msg;
    }

    public T getData() {
        return data;
    }

    public void setData(T data) {
        this.data = data;
    }
}

创建异常类

public class BizExecption extends RuntimeException{
    private int code;
    private String msg;
    public BizExecption(int code,String msg){
        super(msg);
        this.code = code;
        this.msg = msg;
    }

    public int getCode() {
        return code;
    }

    public void setCode(int code) {
        this.code = code;
    }

    public String getMsg() {
        return msg;
    }

    public void setMsg(String msg) {
        this.msg = msg;
    }
}

创建全局异常捕获类，捕获限流异常

@RestControllerAdvice
@Order(-1)
public class GlobleHandlerException {

    /**
     * 捕获限流的异常
     * @param biz
     * @return
     */
    @ExceptionHandler(BizExecption.class)
    public R limiterException(BizExecption biz){
        return new R.Builder().consumerFailed(biz.getCode(),biz.getMsg());
    }
}

创建 aop 类，拦截加了限流注解的控制器

@Aspect
@Slf4j
@Component
public class RateLimiterAop {

    /**
     * 根据请求地址保存不同的令牌桶
     */
    private static final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<String,RateLimiter>();

    /**
     * 切入去点拦截加有  ReqeustRateLimiter 注解的控制器
     */
    @Pointcut("@annotation(RequestRateLimiter)")
    public void pointCut(){}

    @Around("pointCut() && @annotation(requestRateLimiter)")
    public Object doPoint(ProceedingJoinPoint joinPoint,RequestRateLimiter requestRateLimiter) throws Throwable {
        // 获取 request
        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest request = servletRequestAttributes.getRequest();
        // 获取请求 url
        String url = request.getRequestURI();
        // 定义令牌桶
        RateLimiter rateLimiter = null;
        if(!rateLimiterMap. containsKey(url)){
            // 为当前请求创建令牌桶
            rateLimiter = RateLimiter.create(requestRateLimiter.QPS());
            rateLimiterMap.put(url,rateLimiter);
        }
        // 根据请求 url 获取令牌桶
        rateLimiter = rateLimiterMap.get(url);
        // 获取令牌
        boolean acquire = rateLimiter.tryAcquire(requestRateLimiter.acquireTokenTimeout(), requestRateLimiter.timeunit());
        if(acquire){
            // 调用目标方法
            return joinPoint.proceed();
        }
        // 获取不到令牌抛出异常
        throw  new BizExecption(401,requestRateLimiter.resMsg());
    }
}

创建测试类

@RestController
public class RateLimiterController {

    @RequestRateLimiter(QPS = 1,acquireTokenTimeout = 500,timeunit = TimeUnit.MILLISECONDS ,resMsg = "服务器压力有点大呦，稍后再试吧！")
    @RequestMapping("/limiterTest")
    public R limiterTest(){
        return new R.Builder().setDate("获取到令牌请求成功啦！").ok();
    }

    @RequestMapping("/test")
    public R test(){
        return new R.Builder().setDate("不用限流的接口").ok();
    }
}