SpringBoot项目实现自定义注解方式的接口限流

一，实现原理

该限流方式使用的是令牌桶算法，令牌桶算法是基于漏桶算法的一种改进，主要在于令牌桶算法能够在限制服务调用的平均速率的同时，还能够允许一定程度内的突发调用。

1. 系统以固定的速率向桶中添加令牌
2. 当有请求到来时，会尝试从桶中移除一个令牌，如果桶中有足够的令牌，则请求可以被处理或数据包可以被发送；
3. 如果桶中没有令牌，那么请求将被拒绝；
4. 桶中的令牌数不能超过桶的容量，如果新生成的令牌超过了桶的容量，那么新的令牌会被丢弃。
5. 令牌桶算法的一个重要特性是，它能够应对突发流量。当桶中有足够的令牌时，可以一次性处理多个请求，这对于需要处理突发流量的应用场景非常有用。但是又不会无限制的增加处理速率导致压垮服务器，因为桶内令牌数量是有限制的。

如图所示：

二，代码实现

Guava中的RateLimiter就是基于令牌桶实现的，可以直接拿来使用。
1. 引入依赖

		<!--aop切面依赖-->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-aop</artifactId>
        </dependency>
        
        <dependency>
            <groupId>com.google.guava</groupId>
            <artifactId>guava</artifactId>
            <version>33.1.0-jre</version>
        </dependency>

2. 创建注解

/**
 * 限流注解
 */
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimit {

    /**
     * 限流key
     */
    String key() default "";

    /**
     * 限流时间,单位秒
     */
    int time() default 1;

    /**
     * 限流次数
     */
    int count() default 2;

    /**
     * 时间类型
     */
    TimeUnit timeUnit() default TimeUnit.SECONDS;

    /**
     * 提示消息
     */
    String message() default "系统繁忙,请稍后重试";
}

3. AOP切面实现

该切面实现只对接口进行进行限流。

/**
 * 限流切面处理
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {

    private final Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    /**
     * 切面方法，注解之前执行
     */
    @Around("@annotation(rateLimit)")
    public Object doBefore(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
        // key的作用，不同的接口，不同的限流控制
        String key = rateLimit.key();
        RateLimiter rateLimiter;
        if (!rateLimiterMap.containsKey(key)){
            // 创建令牌桶，设置每秒发送得令牌
            rateLimiter = RateLimiter.create(rateLimit.count());
            rateLimiterMap.put(key, rateLimiter);
            log.info("新建了令牌桶={},容量={}", key, rateLimit.count());
        }
        rateLimiter = rateLimiterMap.get(key);
        // 获取令牌，在规定的时间获取令牌，获取不到返回false
        boolean acquire = rateLimiter.tryAcquire(rateLimit.time(), rateLimit.timeUnit());
        // 拿不到令牌，直接返回异常信息
        if (!acquire){
            log.error("令牌桶={},获取令牌失败", key);
            throw new RuntimeException(rateLimit.message());
        }
        return point.proceed();
    }

根据接口参数来区分接口限流的切面实现，参数可以是用户id或者ip地址，这样就实现了具体用户或ip限流

/**
 * 限流切面处理
 */
@Slf4j
@Aspect
@Component
public class RateLimiterAspect {

    /**
     * 方法参数解析器
     */
    private final ParameterNameDiscoverer pnd = new DefaultParameterNameDiscoverer();


    private final Map<String, RateLimiter> rateLimiterMap = Maps.newConcurrentMap();

    /**
     * 切面方法，注解之前执行
     */
    @Around("@annotation(rateLimit)")
    public Object doBefore(ProceedingJoinPoint point, RateLimit rateLimit) throws Throwable {
        // 获取方法（通过方法签名来获取）
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        // 获取参数值
        Object[] args = point.getArgs();
        // 获取方法上参数的名称
        String[] parameterNames = pnd.getParameterNames(method);
        String parameter = "";
        for (int i = 0; i < parameterNames.length; i++) {
            String parameterName = parameterNames[i];
            if (parameterName.equals("param")){
                parameter = (String) args[i];
            }
        }
        // key的作用，不同的接口，不同的流量控制
        String key = rateLimit.key() + parameter ;
        RateLimiter rateLimiter;
        if (!rateLimiterMap.containsKey(key)){
            // 创建令牌桶，设置每秒发送得令牌
            rateLimiter = RateLimiter.create(rateLimit.count());
            rateLimiterMap.put(key, rateLimiter);
            log.info("新建了令牌桶={},容量={}", key, rateLimit.count());
        }
        rateLimiter = rateLimiterMap.get(key);
        // 获取令牌，在规定的时间获取令牌，获取不到返回false
        boolean acquire = rateLimiter.tryAcquire(rateLimit.time(), rateLimit.timeUnit());
        // 拿不到令牌，直接返回异常信息
        if (!acquire){
            log.error("令牌桶={},获取令牌失败", key);
            throw new RuntimeException(rateLimit.message());
        }
        return point.proceed();
    }

3. 注解应用

    @RateLimit(key = "index", count = 5)
    @GetMapping("/index/{param}")
    public String index(@PathVariable("param") String param){
        return param + "hello world!";
    }