自定义注解采用RateLimiter实现限流

首先是自定义注解的实现

//表示该注解被用于方法上
@Target(ElementType.METHOD)
//表示该注解保留到运行时
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimit {
    /**
     * 该注解实际作用，限流
     * rate 代表每秒钟多少请求
     *
     * @return
     */
    int rate() default 100;

    /**
     * 平滑突发限流（SmoothBursty）和平滑预热限流（SmoothWarmingUp）
     * 平滑突发限流，允许突发流量，后面请求速率控制会趋于平稳
     * 平滑预热限流，开始先于小于最大速率执行，慢慢趋于平稳，接近最大速率
     *
     * @return
     */
    String mode() default RateLimitMode.SMOOTH_BURSTY;
}

自定义限流类型

public class RateLimitMode {
    /**
     * 平滑突发限流（SmoothBursty）和平滑预热限流（SmoothWarmingUp）
     * 平滑突发限流，允许突发流量，后面请求速率控制会趋于平稳
     * 平滑预热限流，开始先于小于最大速率执行，慢慢趋于平稳，接近最大速率
     *
     * @return
     */
    public static final String SMOOTH_BURSTY = "smooth-bursty";
    public static final String SMOOTH_WARMING_UP = "smooth-warming-up";

}

切面的实现

@Aspect
@Component
@Order(100)
public class RateLimitAspect {
    /**
     * Guava RateLimiter提供的令牌桶算法可用于平滑突发限流（SmoothBursty）和平滑预热限流（SmoothWarmingUp）实现。
     * 设计思路----------------------------
     * 进入切面后，先进入环绕通知，环绕通知里取出注解限定的速度，并以该速度生产令牌放入桶中，(令牌生产速度为 每秒多少个)，在ConcurrentHashMap中以类名+方法名作为key,以RateLimit作为value
     * * before执行时先去桶中获取令牌，获取不到阻塞一段时间。
     */
    private static Logger LOGGER = LoggerFactory.getLogger(RateLimitAspect.class);
    private static Map<String, Object> RateLimiterMap = new ConcurrentHashMap<String, Object>();
    private volatile static int initServerNum = 0;

    /**
     * 切入点
     */
    @Pointcut("@annotation(注解类全路径)")
    public void rateLimitCutPoint() {
        LOGGER.info("编入切入点");
    }

    @Before("rateLimitCutPoint()")
    public void rateLimitBefore(JoinPoint joinPoint) {
        //获取令牌并计算等待时间
        String methodName = joinPoint.getTarget().getClass().getSimpleName() + "." + joinPoint.getSignature().getName();
        RateLimiter rateLimiter = (RateLimiter) RateLimiterMap.getOrDefault(methodName, RateLimiter.create(134));
        double waitTime = rateLimiter.acquire();
        LOGGER.info("acquire spend {} seconds", waitTime);
    }

    /**
     * 环绕通知必须有返回值, 返回值即为目标方法的返回值
     *
     * @param proceedingJoinPoint
     * @throws Throwable
     */
    @Around("rateLimitCutPoint()")
    public Object rateLimitAround(ProceedingJoinPoint proceedingJoinPoint) throws Throwable {
        try {
            Signature signature = proceedingJoinPoint.getSignature();
            //返回要进行aop方法的全名称
            String methodName = proceedingJoinPoint.getTarget().getClass().getSimpleName() + "." + signature.getName();
            RateLimit rateLimitAnnotation = ((MethodSignature) signature).getMethod().getDeclaredAnnotation(RateLimit.class);
            //注解上的速率
            int rate = rateLimitAnnotation.rate();
            String mode = rateLimitAnnotation.mode();
            putRate(methodName, rate, mode);
            LOGGER.info("限流方法名{},执行速率{}", methodName, ((RateLimiter) RateLimiterMap.get(methodName)).getRate());
        } catch (Exception e) {
            LOGGER.error("限流切面出现异常,异常提示{}", e.getMessage());
        }
        return proceedingJoinPoint.proceed();
    }

    private void putRate(String methodName, int rate, String mode) {
        if (Objects.isNull(RateLimiterMap.get(methodName))) {
            if (mode.equals(RateLimitMode.SMOOTH_BURSTY)) {
                //SmoothBursty模式下,未使用RateLimit最多保存1秒的令牌,即maxBurstSeconds为1秒
                //最大桶数量为maxBurstSeconds 秒生成的令牌数
                RateLimiterMap.putIfAbsent(methodName, RateLimiter.create(rate));
            } else {
                //SmoothWarmingUp模式, permitsPerSecond表示每秒钟新增的令牌数，
                // warmupPeriod表示从冷启动速率过渡到平均速率所需要的时间间隔,设置为1分钟
                RateLimiterMap.putIfAbsent(methodName, RateLimiter.create(rate, 1, TimeUnit.MINUTES));
            }
        }
    }