使用自定义注解+AOP实现RateLimiter令牌桶限流

Springboot 2.x + AOP + RateLimiter，通过添加自定义注解，对请求方法做限流控制。

1. 添加依赖

        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>28.0-jre</version>
        </dependency>

2. 自定义注解

通过注解指定参数，可以自定义限流策略。
自定义RateLimit和guava的RateLimiter有点像，注意区别。

@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {

    /**
     * 每秒创建令牌个数，默认为10
     * @return
     */
    double permitsPerSecond() default 10D;

    /**
     * 获取令牌超时时间
     * @return
     */
    long timeout() default 0;
    
	/**
     * 超时时间单位
     * @return
     */
    TimeUnit timeUnit() default TimeUnit.SECONDS;

}

3. AOP拦截实现

定义切入点@Pointcut，拦截注解@RateLimit 的方法，@Around在目标方法之前织入增强，只有获取到令牌才能执行目标方法。

@Aspect
@Component
public class RateLimiterInterceptor {

    /**
     * 不同的方法存放不同的令牌桶
     */
    private final Map<String, RateLimiter> map = new ConcurrentHashMap<>();

    /**
     * 定义切入点，自定义RateLimit
     */
    @Pointcut("@annotation(com.coco.annotation.RateLimit)")
    public void pointCut() {}

    @Around(value = "pointCut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        System.out.println("进入限流控制");
        Object obj = null;
        JSONObject jsonResult = new JSONObject();
        Signature signature = joinPoint.getSignature();
        String methodName = signature.getName();
        Method method = ((MethodSignature) signature).getMethod();
        //获取注解对象
        RateLimit rateLimit = method.getAnnotation(RateLimit.class);
        if (rateLimit != null) {
            double permitsPerSecond = rateLimit.permitsPerSecond();
            long timeout = rateLimit.timeout();
            TimeUnit timeUnit = rateLimit.timeUnit();
            RateLimiter rateLimiter = null;
            if (!map.containsKey(methodName)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(permitsPerSecond);
                map.put(methodName, rateLimiter);
            }
            rateLimiter = map.get(methodName);
            if (rateLimiter.tryAcquire(timeout, timeUnit)) {
                System.out.println("获取令牌成功，开始执行");
                obj = joinPoint.proceed();
                return obj;
            } else {
                System.out.println("手速不够");
                jsonResult.put("resultMsg", "手速不够");
                return jsonResult;
            }
        } else {
            System.out.println("未知错误");
            jsonResult.put("resultMsg", "未知错误");
            return jsonResult;
        }
    }

}

4. Controller入口方法

通过添加@RateLimit注解实现限流控制，相关参数可自定义。

    @RequestMapping(value = "/getHero/{heroName}", method = RequestMethod.GET)
    @ApiOperation(value = "查看英雄详情")
    @RateLimit(permitsPerSecond = 10D, timeout = 0, timeUnit = TimeUnit.SECONDS)
    public JSONObject getHero(@PathVariable String heroName) {
        JSONObject jsonResult = new JSONObject();
        jsonResult.put("resultMsg", heroName);
        return jsonResult;
    }

5. 限流测试

说明：项目启动后，开启多条线程发送http请求同时访问添加@RateLimite注解的方法。使用CountdownLatch计数器模拟多线程并发：调用await()方法阻塞当前线程，当计数完成后，唤醒所有线程并发执行。

    public static void main(String[] args) throws InterruptedException {
        ExecutorService exec = Executors.newCachedThreadPool();
        RestTemplate restTemplate = new RestTemplate();
        CountDownLatch countDownLatch = new CountDownLatch(1);
        for (int i = 0; i < 15; i ++){
            Runnable runnable = () -> {
                try {
                    countDownLatch.await();
                    Object obj = restTemplate.getForObject("http://localhost:8762/dota/getHero/coco", Object.class);
                    System.out.println(Thread.currentThread().getName() + "： " + JSONObject.toJSONString(obj));
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            };
            exec.submit(runnable);
        }
        countDownLatch.countDown();
        exec.shutdown();
    }

测试结果：