基于Spring的简单分布式限流Filter

使用redis做滑动窗口

talk is cheap show me the code

代码

限流配置RateLimitProperties

@Configuration
@ConfigurationProperties(prefix = "rateLimit")
@Data
public class RateLimitProperties {

    private boolean enabled;

    private List<RateLimitRule> rules;

    /**
     * 每 {timeOfSecond} 秒允许 {key} 命中 {url}正则 {capacity} 次，超过直接返回 {responseBody}
     */
    @Data
    public static class RateLimitRule {
        /**
         * url支持正则, 会以该url作为限速基准。
         * 必填
         */
        private String url;
        /**
         * url的http method，GET POST PUT DELETE 等
         * 为空拦截所有
         */
        private String method;
        /**
         * 被拦截后的响应体
         */
        private String responseBody;
        /**
         * 从http header 或者http parameter中取相应值做拦截基准
         * 必填
         */
        private List<String> keys;
        /**
         * 必填
         */
        private Integer timeOfSecond;
        /**
         * 必填
         */
        private Integer capacity;
    }
}

限流Filter RateLimitFilter

@Component
@WebFilter(urlPatterns = "/*", filterName = "rateLimitFilter")
@Slf4j
public class RateLimitFilter implements Filter {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;
    @Autowired
    private RateLimitProperties rateLimitProperties;
    /**
     * 频率控制脚本
     * 参数：key，窗口长度(毫秒)，窗口容量，发送时间（时间戳：秒）
     * 返回：0：未超出频率；1：超出频率
     * 注意：缓存一周数据，会出现某个时间段断档；允许最大次数需要>=1
     */
    private static final String RATE_LIMIT_LUA = "local nowSize = redis.call('LLEN', KEYS[1])\n" +
            "local window = tonumber(ARGV[1])\n" +
            "local maxSize = tonumber(ARGV[2])\n" +
            "local nowTime = tonumber(ARGV[3])\n" +
            "if nowSize < maxSize then\n" +
            "    redis.call('LPUSH', KEYS[1], nowTime)\n" +
            "    if nowSize == 0 then\n" +
            "        redis.call(\"EXPIRE\", KEYS[1], 86400)\n" +
            "    end\n" +
            "else\n" +
            "    local earliestTime = redis.call('LINDEX', KEYS[1], -1)\n" +
            "    if nowTime - earliestTime <= window then\n" +
            "        return 1\n" +
            "    else\n" +
            "        redis.call('LPUSH', KEYS[1], nowTime)\n" +
            "        redis.call('LTRIM', KEYS[1], 0, maxSize-1)\n" +
            "    end\n" +
            "end\n" +
            "return 0";
    /**
     * rateLimit:URL:KEY:VALUE
     */
    private static final String KEY_PREFIX = "rateLimit:%s:%s:%s";
    private DefaultRedisScript<Long> rateLimitLuaScript;

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        rateLimitLuaScript = new DefaultRedisScript<>();
        rateLimitLuaScript.setResultType(Long.class);
        rateLimitLuaScript.setScriptText(RATE_LIMIT_LUA);
    }

    @Override
    public void doFilter(ServletRequest servletRequest, ServletResponse servletResponse, FilterChain filterChain) throws IOException, ServletException {
        if (!(servletRequest instanceof HttpServletRequest) || !(servletResponse instanceof HttpServletResponse)
                || !rateLimitProperties.isEnabled()) {
            filterChain.doFilter(servletRequest, servletResponse);
            return;
        }
        try {
            HttpServletRequest request = (HttpServletRequest) servletRequest;
            HttpServletResponse response = (HttpServletResponse) servletResponse;
            RateLimitRule rule = findRule(rateLimitProperties, request.getServletPath(), request.getMethod());
            if (rule != null && CollectionUtils.isNotEmpty(rule.getKeys())) {
                for (String key : rule.getKeys()) {
                    String value = getValueFromHeaderOrParam(request, key);
                    if (StringUtils.isBlank(value)) {
                        continue;
                    }
                    boolean block = this.checkRateLimit(key, value, rule);
                    if (block) {
                        log.warn("{}:{}, path:{} is blocked.", key, value, request.getServletPath());
                        response.setStatus(429);
                        response.setContentType("application/json; charset=utf-8");
                        response.setCharacterEncoding("UTF-8");
                        if (StringUtils.isNotBlank(rule.getResponseBody())) {
                            response.getOutputStream().write(rule.getResponseBody().getBytes(StandardCharsets.UTF_8));
                        }
                        return;
                    }
                }
            }
        } catch (Exception e) {
            log.error("rate limit error.", e);
        }
        filterChain.doFilter(servletRequest, servletResponse);
    }

    private boolean checkRateLimit(String param, String value, RateLimitRule rule) {
        String key = String.format(KEY_PREFIX, rule.getUrl(), param, value);
        if (rule.getCapacity() == null || rule.getTimeOfSecond() == null) {
            return false;
        }
        Long count = stringRedisTemplate.execute(rateLimitLuaScript, Lists.newArrayList(key),
                String.valueOf(rule.getTimeOfSecond() * 1000),
                String.valueOf(rule.getCapacity()),
                String.valueOf(System.currentTimeMillis())
        );
        return count != null && count == 1;
    }

    @Override
    public void destroy() {

    }

    public static String getValueFromHeaderOrParam(HttpServletRequest request, String key) {
        String value = request.getHeader(key);
        if (StringUtils.isBlank(value)) {
            value = request.getParameter(key);
        }
        return value;
    }

    public static RateLimitRule findRule(RateLimitProperties rateLimitProperties, String path, String method) {
        if (rateLimitProperties == null || CollectionUtils.isEmpty(rateLimitProperties.getRules())) {
            return null;
        }
        for (RateLimitRule rule : rateLimitProperties.getRules()) {
            if ((StringUtils.isBlank(rule.getMethod()) || method.equalsIgnoreCase(rule.getMethod()))
                    && path.matches(rule.getUrl())) {
                return rule;
            }
        }
        return null;
    }
}

配置文件

rateLimit.enabled=true
rateLimit.rules[0].url=/.**
rateLimit.rules[0].method=GET
rateLimit.rules[0].keys=user_id
rateLimit.rules[0].responseBody={"msg":"您太快了","code":429}
rateLimit.rules[0].timeOfSecond=1
rateLimit.rules[0].capacity=3