使用Guava和AOP实现单机版限流

ゞ浴缸里的玫瑰 2024-03-23 10:50 66阅读 0赞

## 常见的限流算法 ##

计数器限流、漏桶算法，令牌桶算法、滑动窗口

## 基于Guava工具类实现限流 ##

Google开源工具包Guava提供了限流工具类RateLimiter，该类基于令牌桶算法实现流量限制，使用十分方便，而且十分高效，实现步骤如下：

第一步：引入guava依赖包

<dependency>
        <groupId>com.google.guava</groupId>
        <artifactId>guava</artifactId>
        <version>30.1-jre</version>
    </dependency>

第二步：给接口加上限流逻辑

@Slf4j
    @RestController
    @RequestMapping("/limit")
    public class LimitController {
        
        /**
         * 限流策略 ：1秒钟10个请求
         */
        private final RateLimiter limiter = RateLimiter.create(10.0);
    
        private DateTimeFormatter dtf = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss");
    
        @GetMapping("/test1")
        public String testLimiter() {
        
            //500毫秒内，没拿到令牌，就直接进入服务降级
            boolean tryAcquire = limiter.tryAcquire(500, TimeUnit.MILLISECONDS);
    
            if (!tryAcquire) {
        
                log.warn("进入服务降级，时间{}", LocalDateTime.now().format(dtf));
                return "当前排队人数较多，请稍后再试！";
            }
    
            log.info("获取令牌成功，时间{}", LocalDateTime.now().format(dtf));
            return "请求成功";
        }
    }

## 基于AOP实现接口限流 ##

基于AOP的实现方式也非常简单，实现过程如下：

第一步：加入AOP依赖

<dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-aop</artifactId>
    </dependency>

第二步：自定义限流注解

@Retention(RetentionPolicy.RUNTIME)
    @Target({
        ElementType.METHOD})
    @Documented
    public @interface Limit {
        
        /**
         * 资源的key,唯一
         * 作用：不同的接口，不同的流量控制
         */
        String key() default "";
    
        /**
         * 最多的访问限制次数
         */
        double permitsPerSecond () ;
    
        /**
         * 获取令牌最大等待时间
         */
        long timeout();
    
        /**
         * 获取令牌最大等待时间,单位(例:分钟/秒/毫秒) 默认:毫秒
         */
        TimeUnit timeunit() default TimeUnit.MILLISECONDS;
    
        /**
         * 得不到令牌的提示语
         */
        String msg() default "系统繁忙,请稍后再试.";
    }

第三步：使用AOP切面拦截限流注解

@Slf4j
    @Aspect
    @Component
    public class LimitAop {
        
        /**
         * 不同的接口，不同的流量控制
         * map的key为 Limiter.key
         */
        private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();
    
        @Around("@annotation(com.jianzh5.blog.limit.Limit)")
        public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        
            MethodSignature signature = (MethodSignature) joinPoint.getSignature();
            Method method = signature.getMethod();
            //拿limit的注解
            Limit limit = method.getAnnotation(Limit.class);
            if (limit != null) {
        
                //key作用：不同的接口，不同的流量控制
                String key=limit.key();
                RateLimiter rateLimiter = null;
                //验证缓存是否有命中key
                if (!limitMap.containsKey(key)) {
        
                    // 创建令牌桶
                    rateLimiter = RateLimiter.create(limit.permitsPerSecond());
                    limitMap.put(key, rateLimiter);
                    log.info("新建了令牌桶={}，容量={}",key,limit.permitsPerSecond());
                }
                rateLimiter = limitMap.get(key);
                // 拿令牌
                boolean acquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
                // 拿不到命令，直接返回异常提示
                if (!acquire) {
        
                    log.debug("令牌桶={}，获取令牌失败",key);
                    this.responseFail(limit.msg());
                    return null;
                }
            }
            return joinPoint.proceed();
        }
    
        /**
         * 直接向前端抛出异常
         * @param msg 提示信息
         */
        private void responseFail(String msg)  {
        
            HttpServletResponse response=((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getResponse();
            ResultData<Object> resultData = ResultData.fail(ReturnCode.LIMIT_ERROR.getCode(), msg);
            WebUtils.writeJson(response,resultData);
        }
    }

第四步：给需要限流的接口加上注解

@Slf4j
    @RestController
    @RequestMapping("/limit")
    public class LimitController {
        
        
        @GetMapping("/test2")
        @Limit(key = "limit2", permitsPerSecond = 1, timeout = 500, timeunit = TimeUnit.MILLISECONDS,msg = "当前排队人数较多，请稍后再试！")
        public String limit2() {
        
            log.info("令牌桶limit2获取令牌成功");
            return "ok";
        }
    
    
        @GetMapping("/test3")
        @Limit(key = "limit3", permitsPerSecond = 2, timeout = 500, timeunit = TimeUnit.MILLISECONDS,msg = "系统繁忙，请稍后再试！")
        public String limit3() {
        
            log.info("令牌桶limit3获取令牌成功");
            return "ok";
        }
    }