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

喜欢ヅ旅行 2024-04-19 16:09 65阅读 0赞

## 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) {
      ...
    }