利用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<>(); }