利用guava的RateLimiter实现Java接口限流

利用guava的RateLimiter实现Java接口限流

RateLimiter的接口限流方案很简单，主要来说分为三步

• 自定义拦截器

• 注册拦截器

• 自定义注解（可选）

  导入Maven依赖

  <dependency>
      <groupId>com.google.guava</groupId>
      <artifactId>guava</artifactId>
      <version>27.1-jre</version>
  </dependency>
      
  <!-- 阿里JSON解析器 -->
  <dependency>
      <groupId>com.alibaba</groupId>
      <artifactId>fastjson</artifactId>
      <version>1.2.75</version>
  </dependency>
  

  自定义拦截器的方式比较简单，只需要实现HandlerInterceptor接口的preHandle方法就行

  
  @Component
  @Slf4j
  public class RequestLimitingInterceptor implements HandlerInterceptor {
  
      // 创建一个能够分发100个令牌的限速器
      private final static RateLimiter rateLimiter = RateLimiter.create(100);
  
  
      @Override
      public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
          //定义返回信息
          JSONObject jsonObject = new JSONObject();
          jsonObject.put("code", "msg");
          try {
              if (handler instanceof HandlerMethod) {
                  // 获取令牌
                  boolean acquire = rateLimiter.tryAcquire(200, TimeUnit.MILLISECONDS);
                  if (acquire) {
                      //获取令牌成功
                      return true;
                  } else {
                      log.warn("请求被限流,url:{}", request.getServletPath());
                      response(response, toJsonObject(jsonObject));
                      return false;
                  }
              }
              return true;
          } catch (Exception e) {
              e.printStackTrace();
              response(response, toJsonObject(jsonObject));
              return false;
          }
      }
  
      private void response(HttpServletResponse response, JSONObject jo) {
          response.setContentType("application/json; charset=utf-8");
          response.setCharacterEncoding("UTF-8");
          // 语法糖，自动关闭流
          try (PrintWriter out = response.getWriter()) {
              out.append(jo.toJSONString());
          } catch (Exception e) {
              e.printStackTrace();
          }
      }
  
      private JSONObject toJsonObject(Object o) {
          return JSONObject.parseObject(JSON.toJSONString(o));
      }
  
  
  }
  

  注册拦截器，继承WebMvcConfigurationSupport类，重写addInterceptors

  
  @Configuration
  public class WebMvcConfig extends WebMvcConfigurationSupport {
  
      /**
       * 请求限流拦截器
       */
      @Autowired
      protected RequestLimitingInterceptor requestLimitingInterceptor;
  
      @Override
      public void addInterceptors(InterceptorRegistry registry) {
          // 请求限流 根据后面的通配符路径拦截
          registry.addInterceptor(requestLimitingInterceptor).addPathPatterns("/**");
      }
  
  }
  

  自定义注解可选的原因在于，如果需要限制全部接口或者是某一部分接口，那么可以在拦截器注册时通过通配符路径完成对他们的限制，这样可以避免在每个接口上面加注解的麻烦事，如果想实现不同接口有不同的限制效果，那么就需要用到自定义注解

  import java.lang.annotation.*;
  import java.util.concurrent.TimeUnit;
   
  /**
   * 限速注解
   */
  @Target({ElementType.METHOD})
  @Retention(RetentionPolicy.RUNTIME)
  @Documented
  public @interface Limiter {
      /**
       * 每秒创建令牌个数，默认:200
       */
      double QPS() default 200D;
   
      /**
       * 获取令牌等待超时时间 默认:500
       */
      long timeout() default 500;
   
      /**
       * 超时时间单位 默认:毫秒
       */
      TimeUnit timeunit() default TimeUnit.MILLISECONDS;
   
      /**
       * 限速提示信息
       */
      String msg() default "请稍后再试!";
  }
  

  使用注解的话需要同步修改拦截器的实现

  @Component
  @Slf4j
  public class RequestLimitingInterceptor implements HandlerInterceptor {
      private final Map<String, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();
   
      @Override
      public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
          //定义返回信息
          JSONObject jsonObject = new JSONObject();
          jsonObject.put("code", "msg");
          try {
              if (handler instanceof HandlerMethod) {
                  HandlerMethod handlerMethod = (HandlerMethod) handler;
                  Limiter rateLimit = handlerMethod.getMethodAnnotation(Limiter.class);
                  //判断是否有注解
                  if (rateLimit != null) {
                      // 获取请求url
                      String url = request.getRequestURI();
                      RateLimiter rateLimiter;
                      // 判断map集合中是否有创建好的令牌桶
                      if (!rateLimiterMap.containsKey(url)) {
                          // 创建令牌桶
                          rateLimiter = RateLimiter.create(rateLimit.QPS());
                          rateLimiterMap.put(url, rateLimiter);
                      }
                      rateLimiter = rateLimiterMap.get(url);
                      // 获取令牌
                      boolean acquire = rateLimiter.tryAcquire(rateLimit.timeout(), rateLimit.timeunit());
                      if (acquire) {
                          //获取令牌成功
                          return true;
                      } else {
                          log.warn("请求被限流,url:{}", request.getServletPath());
                          response(response, toJsonObject(jsonObject));
                          return false;
                      }
                  }
              }
              return true;
          } catch (Exception e) {
              e.printStackTrace();
              response(response, toJsonObject(jsonObject));
              return false;
          }
      }
   
      private void response(HttpServletResponse response, JSONObject jo) {
          response.setContentType("application/json; charset=utf-8");
          response.setCharacterEncoding("UTF-8");
          // 语法糖，自动关闭流
          try (PrintWriter out = response.getWriter()) {
              out.append(jo.toJSONString());
          } catch (Exception e) {
              e.printStackTrace();
          }
      }
   
      private JSONObject toJsonObject(Object o) {
          return JSONObject.parseObject(JSON.toJSONString(o));
      }
  

  注解使用姿势

  @RequestLimiter(QPS = 5D, timeout = 200, timeunit = TimeUnit.MILLISECONDS,msg = "玩命加载中,请稍后再试")
  @PostMapping("/list")
  public String list(){
        return new ArrayList<>();
  }