基于Guava RateLimiter，实现一个“API级别的限流”注解

首先，解释一下标题。
我们有一个需求：能为每个接口单独设置一个限流值。那么每个接口都需要增加相应的代码，只有自己写一个注解，使用成本才低，对业务代码的侵入也低。

---

一、整体思路

1. 自定义一个注解，里面有个限流值的变量；
2. 在需要的接口上，加上该注解，并设置好限流值，比如：@RateLimit(5)；
3. 写一个针对该注解的切面，before()阶段进行限流判断和限流处理。

二、开始编写代码

1. 自定义注解

@Inherited
@Documented
@Target({ElementType.METHOD}) //方法级别
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
     /**
     * 为了方便管理各个接口的限流值，这里是设置对应application.properties里的属性key，key对应的value才是限流值
     */
    String limitKey() default "api.default.limit";
}

2. 每个API上使用该注解

@ApiOperation(value = "查询产品详情")
@RequestMapping(value = "/{productId}", method = RequestMethod.GET)
@RateLimit(limitKey = "xxxxx.xxxx")
public Result queryProduct(@PathVariable("productId") Integer productId) {
       return new Result<>(xxClient.queryProduct(productId));
}

3. 针对该注解的切面（限流的核心）

@Component
@Aspect
@Slf4j
public class RateLimitAspect {
    /**
     * 针对含有该注解的地方进行切入
     */
    @Pointcut("@annotation(com.xxx.xxx.RateLimit)")
    public void serviceLimit() {
    }
    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
            //限流判断，和限流处理
    }

切面的大体框架有了，我们现在就剩具体的限流逻辑代码

三、Guava RateLimiter

限流有很多算法，包括令牌桶、漏桶。根据需求，我们选择的是令牌桶。

令牌桶：系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则拒绝服务。

而令牌桶的具体实现，如果我们自己完成，就有些复杂了，所以这里我选择了Guava依赖包里的RateLimiter。先看代码吧

        
            com.google.guava
            guava
            20.0
        

@Autowired
    private Environment env;

    //用来存放不同接口的RateLimiter(key为接口名称，value为RateLimiter)
    private ConcurrentHashMap rateLimitMap = new ConcurrentHashMap<>();

    @Around("serviceLimit()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        Object obj = null;
        //获取拦截的签名
        Signature sig = joinPoint.getSignature();
        //获取拦截的方法名
        MethodSignature msig = (MethodSignature) sig;
        //返回被织入增加处理目标对象
        Object target = joinPoint.getTarget();
        //为了获取注解信息
        Method currentMethod = target.getClass().getMethod(msig.getName(), msig.getParameterTypes());
        //获取注解信息
        RateLimit annotation = currentMethod.getAnnotation(RateLimit.class);
        //在application.properties中获取注解每秒加入桶中的token
        String limitNum = env.getProperty(annotation.limitKey());
        // 注解所在方法名区分不同的限流策略
        String functionName = msig.toShortString();
        //如果值为0，则为不进行限流
        if (limitNum == null || limitNum.equals("0")) {
            obj = joinPoint.proceed();
        } else {
            //获取rateLimiter
            if(!rateLimitMap.containsKey(functionName)){
                //使用最简洁的方法来创建RateLimiter，RateLimiter.create(double xx)，如果有需要，可自行设置RateLimiter其他属性
                rateLimitMap.put(functionName, RateLimiter.create(Double.parseDouble(limitNum)));
            }
            RateLimiter rateLimiter = rateLimitMap.get(functionName);
                        //尝试获得一个令牌
            if (rateLimiter.tryAcquire(1)) {
                //执行方法
                obj = joinPoint.proceed();
            } else {
                //拒绝了请求（服务降级），这是自己定义的异常类
                throw new CommonException("拒绝了访问open api方法的请求", FieldCodeEnum.OPEN_API, null, DetailCodeEnum.ACCESS, ErrorCodeEnum.OUT_OF_LIMIT);
            }
        }
        return obj;
    }

api.default.limit=5

四、完成

可以使用PostMan软件发起压力测试请求，比如1秒钟发出去10个，结果会是6个通行了，但后4个被拒绝了。
明明我们设置的限流值是5，为何6个通行了呢？见官方描述：

When the incoming request rate exceeds {@code permitsPerSecond} the
rate limiter will release one permit every {@code (1.0 / permitsPerSecond)} seconds.

这就是 RateLimiter 的预加载功能。

---

以上只是我搭建了一个限流的基础框架，大家可以继续完善它，以实现自己的需要。
欢迎大家一起交流相关知识