使用 RateLimiter 令牌桶算法实现自定义注解限流（修改桶的容量）

1、创建注解类 定义相关参数

// 允许注解加在方法上
@Target(ElementType.METHOD)
// 允许运行时获取
@Retention(RetentionPolicy.RUNTIME)
public @interface ApiLimit {

    /**
     * 接口名称
     * @return
     */
    String name() default "";

    /**
     * 每秒生成令牌数 默认0.2既5秒允许调用一次
     * @return
     */
    double timesPerSecond() default 0.2;

    /**
     * 未获取到令牌允许等待时长（单位秒） 默认等待10秒
     * @return
     */
    long waitingTime() default 10;
	
	/**
     * 桶的容量最大存储多少秒-因为RateLimiter 的容量存储秒数是私有的，
     * 并且默认是1秒不提供外部设置，所以要改变容量，我这边采用的是反射来修改它的私有变量
     * @return
     */
    long maxBurstSeconds() default 40;
}

2、创建aop相关类 利用aop的环绕通知实现拦截 统一获取令牌后放行

@Slf4j
@Aspect
@Component
public class ApiLimitAspect {

    /**
     * 定义一个线程安全的map集合
     */
    private ConcurrentHashMap<String, RateLimiter> rateLimiters = new ConcurrentHashMap();

    /**
     * 环绕通知 - 监控带有ApiLimit注解的方法
     * @param joinPoint
     * @return
     */
    @Around(value = "@annotation(com.rensenergy.annotation.ApiLimit)")
    public Object around(ProceedingJoinPoint joinPoint) {
        try {
            //获取拦截的方法名
            Signature sig = joinPoint.getSignature();
            MethodSignature methodSignature = (MethodSignature) sig;
            RequestAttributes ra = RequestContextHolder.getRequestAttributes();
            ServletRequestAttributes sra = (ServletRequestAttributes) ra;

            HttpServletRequest request = sra.getRequest();
            // 获取请求的ip
            String requestIp = SecurityUtils.getIp(request);
            // 请求参数
            String requestParams = JacksonUtils.serialize(joinPoint.getArgs(),true);
            log.info("方法名："+methodSignature.getName()+"===============请求参数："+requestParams);
            // 获取方法上的注解信息
            ApiLimit apiLimit = methodSignature.getMethod().getDeclaredAnnotation(ApiLimit.class);
            if (apiLimit == null) {
                // 执行目标方法
                return joinPoint.proceed();
            }
            // ip 加上 注解上的name 组成限流name
            String name = requestIp + apiLimit.name();
            // 获取注解上的timesPerSecond
            double timesPerSecond = apiLimit.timesPerSecond();
            // 获取注解上的等待时间
            long waitingTime = apiLimit.waitingTime();
            // 获取注解上的存储时长
            long maxBurstSeconds = apiLimit.maxBurstSeconds();
            // 从集合中获取rateLimiter
            RateLimiter rateLimiter = rateLimiters.get(name);
            if (rateLimiter == null) {
                // 没有则创建
                // 没有则创建
                rateLimiter = RateLimiter.create(timesPerSecond);
                Class calss = rateLimiter.getClass();
                // 利用反射修改maxBurstSeconds 值
                Field field = calss.getDeclaredField("maxBurstSeconds");
                field.setAccessible(true);
                field.set(rateLimiter, maxBurstSeconds);
                rateLimiter.setRate(timesPerSecond);
                rateLimiters.put(name, rateLimiter);
            }
            Object result ;
            // 等待获取令牌
            if(rateLimiter.tryAcquire(waitingTime, TimeUnit.SECONDS)){
                // 在等待时间内获取到令牌就执行方法
                result = joinPoint.proceed();
            }else{
                // 未获取到就返回系统繁忙
                result = ResponseHelper.failed(ResultCode.SERVICE_BUSY);
            }
            return result;
        } catch (Throwable throwable) {
            throw new RuntimeException("服务器异常！");
        }
    }
}

3、在方法上带上注解

@ApiLimit(name = "limitTest", timesPerSecond = 1, waitingTime = 12)
@PostMapping("/limitTest")
public ResponseEntity<?> passBack(@RequestBody JSONObject jsonObject){
    String str = jsonObject.toJSONString();
    log.info("======================="+str);
    return new ResponseEntity<>(0, "成功", null);
}

Guava令牌桶算法原理参考 Guava令牌桶源码分析