cn.yueshutong.springbootstartercurrentlimiting.interceptor.DefaultCurrentInterceptor Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of spring-boot-starter-current-limiting Show documentation

Current limiting tool based on leak bucket algorithm and token bucket algorithm.

There is a newer version: 0.1.1.RELEASE

Show newest version

package cn.yueshutong.springbootstartercurrentlimiting.interceptor;

import cn.yueshutong.springbootstartercurrentlimiting.common.SpringContextUtil;
import cn.yueshutong.springbootstartercurrentlimiting.core.RateLimiter;
import cn.yueshutong.springbootstartercurrentlimiting.core.RateLimiterCloud;
import cn.yueshutong.springbootstartercurrentlimiting.core.RateLimiterSingle;
import cn.yueshutong.springbootstartercurrentlimiting.handler.CurrentInterceptorHandler;
import cn.yueshutong.springbootstartercurrentlimiting.properties.CurrentProperties;
import cn.yueshutong.springbootstartercurrentlimiting.properties.CurrentRuleProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * 系统默认限流规则：通过配置文件指定参数
 * Create by [email protected] 2019/4/21 0021 11:19
 */
public class DefaultCurrentInterceptor implements HandlerInterceptor {
    private RateLimiter rateLimiter;
    private CurrentRuleProperties limiterRule;
    private CurrentInterceptorHandler interceptorHandler;
    private Logger logger = LoggerFactory.getLogger(getClass());

    DefaultCurrentInterceptor(CurrentProperties properties, CurrentInterceptorHandler handler, CurrentRuleProperties limiterRule) {
        this.limiterRule = limiterRule;
        this.interceptorHandler = handler;
        initRateLimiter(properties, limiterRule);
    }

    private void initRateLimiter(CurrentProperties properties, CurrentRuleProperties limiterRule) {
        if (properties.isCloudEnabled()) {
            this.rateLimiter = RateLimiterCloud.of(limiterRule.getQps(), limiterRule.getInitialDelay(), SpringContextUtil.getApplicationName(), limiterRule.isOverflow());
        } else {
            this.rateLimiter = RateLimiterSingle.of(limiterRule.getQps(), limiterRule.getInitialDelay(), limiterRule.isOverflow());
        }
    }

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (limiterRule.isFailFast()) { //执行快速失败
            return tryAcquireFailed(request, response);
        } else { //执行阻塞策略
            return rateLimiter.tryAcquire();
        }
    }

    private boolean tryAcquireFailed(HttpServletRequest request, HttpServletResponse response) throws Exception {
        if (rateLimiter.tryAcquireFailed()) { //取到令牌
            return true;
        } else { //没取到令牌
            if (interceptorHandler == null) {
                response.getWriter().print(RateLimiter.message);
            } else {
                interceptorHandler.preHandle(request, response);
            }
            return false;
        }
    }

    @Override
    public void postHandle(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, ModelAndView modelAndView) {
        //视图渲染之前
    }

    @Override
    public void afterCompletion(HttpServletRequest httpServletRequest, HttpServletResponse httpServletResponse, Object o, Exception e) {
        //整个请求结束之后
    }
}