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

令牌桶

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

令牌桶思想

大小固定的令牌桶可自行以恒定的速率源源不断地产生令牌。如果令牌不被消耗,或者被消耗的速度小于产生的速度,令牌就会不断地增多,直到把桶填满。

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

限流demo实现思路
  • 创建令牌桶类
  • 项目启动初始化令牌桶,并设置定时器,定时向桶内放入令牌
  • 自定义限流注解,在需要限流的接口上打上注解
  • 配置令牌桶拦截器,对所有路径进行拦截,对无限流注解的接口直接放行,对有限流注解的做取令牌处理,取到令牌即放行,没取到令牌即抛出自定义异常
  • 自定义异常并使用AOP做全局异常处理

 

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

BucketUtil如下

 1 public class BucketUtil {
 2 
 3     //默认容量10
 4     static final int DEFAULT_MAX_COUNT = 10;
 5     // 默认增长速率为1
 6     static final int DEFAULT_CREATE_RATE = 1;
 7     // 使用HashMap存放令牌桶,这里默认为10个令牌桶
 8     public static  HashMap<String, BucketUtil> buckets = new HashMap(10);
 9 
10     //自定义容量,一旦创建不可改变
11     final int maxCount;
12     //自定义增长速率1s几个令牌
13     int createRate;
14     //当前令牌数
15     int size=0;
16 
17 
18 
19     // 默认令牌桶的容量及增长速率
20     public BucketUtil() {
21         maxCount = DEFAULT_MAX_COUNT;
22         createRate = DEFAULT_CREATE_RATE;
23     }
24     // 自定义令牌桶容量及增长速率
25     public BucketUtil(int maxCount, int createRate) {
26         this.maxCount = maxCount;
27         this.createRate = createRate;
28     }
29 
30     public int getSize() {
31         return size;
32     }
33 
34     public boolean isFull() {
35         return size == maxCount;
36     }
37 
38     //根据速率自增生成一个令牌
39     public synchronized void incrTokens() {
40         for (int i = 0; i < createRate; i++)
41         {
42             if (isFull())
43                 return;
44             size++;
45         }
46     }
47 
48     // 取一个令牌
49     public synchronized boolean getToken() {
50         if (size > 0)
51             size--;
52         else
53             return false;
54         return true;
55     }
56 
57     @Override
58     public boolean equals(Object obj) {
59         if (obj == null)
60             return false;
61         BucketUtil bucket = (BucketUtil) obj;
62         if (bucket.size != size || bucket.createRate != createRate || bucket.maxCount != maxCount)
63             return false;
64         return true;
65     }
66 
67     @Override
68     public int hashCode() {
69         return Objects.hash(maxCount, size, createRate);
70     }
71 
72 }
BucketUtil

初始化令牌桶

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

 1 @EnableScheduling
 2 @SpringBootApplication
 3 public class DemoApplication {
 4 
 5  public static void main(String[] args) {
 6   SpringApplication.run(DemoApplication.class, args);
 7   // 为了方便测试这里定义1容量  1增长速率
 8   BucketUtil bucketUtil = new BucketUtil(1,1);
 9   // 生成名为:bucket的令牌桶
10   BucketUtil.buckets.put("bucket",bucketUtil);
11  }
12  @Scheduled(fixedRate = 1000)// 定时1s
13  public void timer() {
14   if (BucketUtil.buckets.containsKey("bucket")){
15    //名为:bucket的令牌桶 开始不断生成令牌
16    BucketUtil.buckets.get("bucket").incrTokens();
17   }
18  }
19 }
DemoApplication

自定义注解以及异常

@Target({ElementType.METHOD})// METHOD代表是用在方法上
@Retention(RetentionPolicy.RUNTIME)
public @interface BucketAnnotation {
}
BucketAnnotation
1 public class APIException extends RuntimeException {
2     private static final long serialVersionUID = 1L;
3     private String msg;
4     public APIException(String msg) {
5         super(msg);
6         this.msg = msg;
7     }
8 }
APIException

配置拦截器

 1 /**
 2  * 令牌桶拦截器
 3  */
 4 public class BucketInterceptor implements HandlerInterceptor {
 5 
 6     // 预处理回调方法,在接口调用之前使用  true代表放行  false代表不放行
 7     @Override
 8     public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
 9         if (!(handler instanceof HandlerMethod)) {
10             return true;
11         }
12 
13         HandlerMethod handlerMethod = (HandlerMethod) handler;
14         Method method = handlerMethod.getMethod();
15 
16         BucketAnnotation methodAnnotation = method.getAnnotation(BucketAnnotation.class);
17         if (methodAnnotation!=null){
18          // 在名为:bucket的令牌桶里取令牌 取到即放行 未取到即抛出异常
19             if(BucketUtil.buckets.get("bucket").getToken()){
20                 return true;
21             }
22             else{
23              // 抛出自定义异常
24                 throw new APIException("不好意思,您被限流了");
25             }
26         }else {
27             return true;
28         }
29     }
30     // 接口调用之后,返回之前 使用
31     @Override
32     public void postHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, ModelAndView modelAndView) throws Exception {
33     }
34 
35     // 整个请求完成后,在视图渲染前使用
36     @Override
37     public void afterCompletion(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, Exception e) throws Exception {
38     }
39 }
BucketInterceptor

将拦截器注入

 1 @Configuration
 2 public class WebMvcConfg implements WebMvcConfigurer {
 3 
 4     @Override
 5     public void addInterceptors(InterceptorRegistry registry) {
 6         // 令牌桶拦截器 添加拦截器并选择拦截路径
 7         registry.addInterceptor(bucketInterceptor()).addPathPatterns("/**");
 8     }
 9     @Bean
10     public BucketInterceptor bucketInterceptor() {
11         return new BucketInterceptor();
12     }
13 }
WebMvcConfg

AOP全局异常处理

1 @RestControllerAdvice
2 public class WebExceptionControl {
3     @ExceptionHandler(APIException.class)
4     public E3Result APIExceptionHandler(APIException e) {
5         return E3Result.build(400,e.getMessage());
6     }
7 }
WebExceptionControl

测试

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

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

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

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

 

 

 

 

总结

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

  • 分布式环境下不适用
  • 令牌桶可以有多个,不同的接口采用不同令牌桶的时候,拦截器无法分开限流
  • 一次请求消耗一个令牌,可以被恶意消耗

改进方法:

  • 令牌桶实现采用redis集群存取
  • 注解添加value参数,可以给对应接口打上对应的令牌桶参数,拦截器需对注解参数校验,实现多个接口多个令牌桶的限流
  • 对用户IP校验限制次数,防止恶意攻击

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


 

转自:https://mp.weixin.qq.com/s/6Uh6e9T93osxttpL6M5cQA

posted on 2022-01-12 18:07  腾逸  阅读(549)  评论(0编辑  收藏  举报