guava ratelimiter 实现对不同接口单独进行动态限流

1. 背景

项目中不同接口的流量限制不一样，需要对不同接口单独进行限流

2.  设计

1） 不同接口的流量峰值注册在配置中心

2） 在拦截器中获取接口的限流配置信息，针对不同的接口单独实例化RateLimiter对象

3） 进行流量操作

3 代码实现

 3.1 拦截器实现

import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.ModelAndView;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;

public abstract class AbstractAnnotationHandlerInterceptorAdapter<T extends Annotation> extends HandlerInterceptorAdapter {
    private static final Logger log = LoggerFactory.getLogger(AbstractAnnotationHandlerInterceptorAdapter.class);
    private static final String ANNOTATION_ATTRIBUTE_NAME = "__Annotation";
    private static final String ANNOTATION_CLASS_ATTRIBUTE_NAME = "__AnnotationClass";

    public AbstractAnnotationHandlerInterceptorAdapter() {
    }

    protected abstract Class<T> getRequiredAnnotationClass();

    protected boolean matchRequiredAnnotation(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        log.trace("Check required annotation start ...");
        if (this.getRequiredAnnotationClass() != null && handler.getClass().isAssignableFrom(HandlerMethod.class)) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            Annotation annotation = AnnotationUtils.findAnnotation(method, this.getRequiredAnnotationClass());
            if (annotation == null) {
                annotation = AnnotationUtils.findAnnotation(method.getDeclaringClass(), this.getRequiredAnnotationClass());
            }

            if (annotation != null) {
                request.setAttribute("__Annotation", annotation);
                request.setAttribute("__AnnotationClass", this.getRequiredAnnotationClass());
                log.trace("Match required annotation : {}", annotation);
                return true;
            }
        }

        return false;
    }

    protected Object before(HttpServletRequest request, HttpServletResponse response, Object handler, T anno) throws Exception {
        return true;
    }

    protected void after(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView, T anno) throws Exception {
    }

    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (this.matchRequiredAnnotation(request, response, handler)) {
            Object result = this.before(request, response, handler, (T) request.getAttribute("__Annotation"));
            if (result != Boolean.TRUE) {
                return false;
            }
        }

        return true;
    }

    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        if (this.getRequiredAnnotationClass() != null && this.getRequiredAnnotationClass() == request.getAttribute("__AnnotationClass")) {
            this.after(request, response, handler, modelAndView, (T) request.getAttribute("__Annotation"));
            request.removeAttribute("__Annotation");
            request.removeAttribute("__AnnotationClass");
        }

    }

    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        super.afterCompletion(request, response, handler, ex);
    }

    public void afterConcurrentHandlingStarted(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        super.afterConcurrentHandlingStarted(request, response, handler);
    }
}

3.2 实现类 

public class RateLimiterInterceptorAdapter extends AbstractAnnotationHandlerInterceptorAdapter<RateLimiterValidate> {

    @Override
    protected Class<RateLimiterValidate> getRequiredAnnotationClass() {
        return RateLimiterValidate.class;
    }

    protected Object before(HttpServletRequest request, HttpServletResponse response, Object handler, RateLimiterValidate validate) throws Exception {
        String key = request.getServletPath();
        log.info("RateLimiterCache checkAndAdd start, path {}", key);
        try {
            Map<String, Integer> limitedPathMap = QuickAppConfigurationManager.getQuickAppConfiguration().getRateLimitMap();
            if (limitedPathMap.containsKey(key)) {
                // 判断是否超过限制
                if (RateLimiterCache.checkAndGet(key, limitedPathMap.get(key)).tryAcquire()) {
                    log.error("check limit pass, path {}", key);
                    return true;
                }
                log.error("check limit not pass, path {}", key);
                // 超过限制，直接返回提示信息
                handlerRateLimitedResponse(key, response);
                return false;
            }
        } catch (Exception e) {
            log.error("RateLimiterCache checkAndAdd error, path {}", request.getServletPath(), e);
        }
        return true;
    }

    private void handlerRateLimitedResponse(String key, HttpServletResponse response) {
        response.setStatus(HttpStatus.SC_OK);
        response.setHeader("Content-Type", "application/json");
        try {
            response.getWriter().print(JSON.toJSONString(CommonResult.<String>build().code(Constant.ErrorCode.EXCEED_LIMIT).message("请求繁忙中,请稍后重试")));
        } catch (IOException e) {
            log.error("handlerRateLimitedResponse error,key {}", key, e);
        }
    }

3.3 注解

@Target({ElementType.METHOD, ElementType.TYPE})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiterValidate {
    @AliasFor("required")
    boolean value() default true;

    @AliasFor("value")
    boolean required() default true;
}

 

 

4 如何使用限流

直接加一个注解

    @PostMapping("/rebind")
    @RateLimiterValidate
    public CommonResult<Boolean> rebind(HttpServletRequest request, AuthOrderDto authOrderDto) {
  ...
}