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 }
初始化令牌桶
在启动类上初始化并生成定时器

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 }
自定义注解以及异常

@Target({ElementType.METHOD})// METHOD代表是用在方法上 @Retention(RetentionPolicy.RUNTIME) public @interface 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 }
配置拦截器

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 }
将拦截器注入

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 }
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 }
测试
在我们需要限流的接口上打上自定义注解,如下

@BucketAnnotation @RequestMapping(value = "/bucket") public E3Result bucket(){ return E3Result.ok("访问成功"); }
关于E3Result只是一个封装好的返回类,这里就不贴出来了,大家有的替换成自己的,没有的可以直接用String型测试
上面为了方便测试,令牌桶的容量设置成了1,所以这是取到令牌成功的
总结
上面的限流只是一个demo还有很多不足的地方,如:
- 分布式环境下不适用
- 令牌桶可以有多个,不同的接口采用不同令牌桶的时候,拦截器无法分开限流
- 一次请求消耗一个令牌,可以被恶意消耗
改进方法:
- 令牌桶实现采用redis集群存取
- 注解添加value参数,可以给对应接口打上对应的令牌桶参数,拦截器需对注解参数校验,实现多个接口多个令牌桶的限流
- 对用户IP校验限制次数,防止恶意攻击
实际项目限流会更加严谨,上述只是提供了一个思路以及演示demo,不喜勿喷谢谢。
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 【.NET】调用本地 Deepseek 模型
· CSnakes vs Python.NET:高效嵌入与灵活互通的跨语言方案对比
· 我与微信审核的“相爱相杀”看个人小程序副业
· DeepSeek “源神”启动!「GitHub 热点速览」
· Plotly.NET 一个为 .NET 打造的强大开源交互式图表库
2017-01-12 Java 集合系列06之 Vector详细介绍(源码解析)和使用示例