SpringBoot使用令牌桶算法+拦截器+自定义注解+自定义异常实现简单的限流

点击上方“Java基基”，选择“设为星标”

做积极的人，而不是积极废人！

每天14:00更新文章，每天掉亿点点头发...

源码精品专栏

在高并发的情况下，限流是后端常用的手段之一内容拦截器有什么用，可以对系统限流、接口限流、用户限流等，本文就使用令牌桶算法+拦截器+自定义注解+自定义异常实现限流的demo。

令牌桶思想

大小固定的令牌桶可自行以恒定的速率源源不断地产生令牌。如果令牌不被消耗内容拦截器有什么用，或者被消耗的速度小于产生的速度，令牌就会不断地增多，直到把桶填满。

后面再产生的令牌就会从桶中溢出。最后桶中可以保存的最大令牌数永远不会超过桶的大小。然后每个访问的用户都会从中取走一块令牌，取到了令牌才能访问，如果没取到令牌即代表已达到访问上限，将被限流不允许访问

“

推荐下自己做的 Spring Boot 的实战项目：

这里为了防止并发问题在生成令牌和取令牌的方法上加了synchronized

BucketUtil如下

public class BucketUtil {

    //默认容量10
    static final int DEFAULT_MAX_COUNT = 10;
    // 默认增长速率为1
    static final int DEFAULT_CREATE_RATE = 1;
    // 使用HashMap存放令牌桶，这里默认为10个令牌桶
    public static  HashMap buckets = new HashMap(10);

    //自定义容量，一旦创建不可改变
    final int maxCount;
    //自定义增长速率1s几个令牌
    int createRate;
    //当前令牌数
    int size=0;



    // 默认令牌桶的容量及增长速率
    public BucketUtil() {
        maxCount = DEFAULT_MAX_COUNT;
        createRate = DEFAULT_CREATE_RATE;
    }
    // 自定义令牌桶容量及增长速率
    public BucketUtil(int maxCount, int createRate) {
        this.maxCount = maxCount;
        this.createRate = createRate;
    }

    public int getSize() {
        return size;
    }

    public boolean isFull() {
        return size == maxCount;
    }

    //根据速率自增生成一个令牌
    public synchronized void incrTokens() {
        for (int i = 0; i < createRate; i++)
        {
            if (isFull())
                return;
            size++;
        }
    }

    // 取一个令牌
    public synchronized boolean getToken() {
        if (size > 0)
            size--;
        else
            return false;
        return true;
    }

    @Override
    public boolean equals(Object obj) {
        if (obj == null)
            return false;
        BucketUtil bucket = (BucketUtil) obj;
        if (bucket.size != size || bucket.createRate != createRate || bucket.maxCount != maxCount)
            return false;
        return true;
    }

    @Override
    public int hashCode() {
        return Objects.hash(maxCount, size, createRate);
    }

}

“

推荐下自己做的 Spring Cloud 的实战项目：

在启动类上初始化并生成定时器

@EnableScheduling
@SpringBootApplication
public class DemoApplication {

 public static void main(String[] args) {
  SpringApplication.run(DemoApplication.class, args);
  // 为了方便测试这里定义1容量  1增长速率
  BucketUtil bucketUtil = new BucketUtil(1,1);
  // 生成名为：bucket的令牌桶
  BucketUtil.buckets.put("bucket",bucketUtil);
 }
 @Scheduled(fixedRate = 1000)// 定时1s
 public void timer() {
  if (BucketUtil.buckets.containsKey("bucket")){
   //名为：bucket的令牌桶 开始不断生成令牌
   BucketUtil.buckets.get("bucket").incrTokens();
  }
 }
}

BucketAnnotation注解

@Target({ElementType.METHOD})// METHOD代表是用在方法上
@Retention(RetentionPolicy.RUNTIME)
public @interface BucketAnnotation {
}

APIException 异常

public class APIException extends RuntimeException {
    private static final long serialVersionUID = 1L;
    private String msg;
    public APIException(String msg) {
        super(msg);
        this.msg = msg;
    }
}

BucketInterceptor拦截器如下

/**
 * 令牌桶拦截器
 */
public class BucketInterceptor implements HandlerInterceptor {

    // 预处理回调方法，在接口调用之前使用  true代表放行  false代表不放行
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }

        HandlerMethod handlerMethod = (HandlerMethod) handler;
        Method method = handlerMethod.getMethod();

        BucketAnnotation methodAnnotation = method.getAnnotation(BucketAnnotation.class);
        if (methodAnnotation!=null){
         // 在名为：bucket的令牌桶里取令牌 取到即放行 未取到即抛出异常
            if(BucketUtil.buckets.get("bucket").getToken()){
                return true;
            }
            else{
             // 抛出自定义异常
                throw new APIException("不好意思，您被限流了");
            }
        }else {
            return true;
        }
    }
    // 接口调用之后，返回之前 使用
    @Override
    public void postHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, ModelAndView modelAndView) throws Exception {
    }

    // 整个请求完成后，在视图渲染前使用
    @Override
    public void afterCompletion(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, Exception e) throws Exception {
    }
}

将拦截器注入

@Configuration
public class WebMvcConfg implements WebMvcConfigurer {

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        // 令牌桶拦截器 添加拦截器并选择拦截路径
        registry.addInterceptor(bucketInterceptor()).addPathPatterns("/**");
    }
    @Bean
    public BucketInterceptor bucketInterceptor() {
        return new BucketInterceptor();
    }
}

@RestControllerAdvice
public class WebExceptionControl {
    @ExceptionHandler(APIException.class)
    public E3Result APIExceptionHandler(APIException e) {
        return E3Result.build(400,e.getMessage());
    }
}

在我们需要限流的接口上打上自定义注解，如下

@BucketAnnotation
@RequestMapping(value = "/bucket")
public  E3Result bucket(){
    return E3Result.ok("访问成功");
}

关于E3Result只是一个封装好的返回类，这里就不贴出来了，大家有的替换成自己的，没有的可以直接用String型测试

上面为了方便测试，令牌桶的容量设置成了1，所以这是取到令牌成功的

这是没取到令牌被限流的

上面的限流只是一个demo还有很多不足的地方，如：

改进方法：

实际项目限流会更加严谨，上述只是提供了一个思路以及演示demo，不喜勿喷谢谢。

欢迎加入我的知识星球，一起探讨架构，交流源码。加入方式，长按下方二维码噢：

已在知识星球更新源码解析如下：

最近更新《芋道 SpringBoot 2.X 入门》系列，已经 101 余篇，覆盖了MyBatis、Redis、MongoDB、ES、分库分表、读写分离、SpringMVC、Webflux、权限、WebSocket、Dubbo、RabbitMQ、RocketMQ、Kafka、性能测试等等内容。

提供近 3W 行代码的 SpringBoot 示例，以及超 6W 行代码的电商微服务项目。