使用redis做滑动窗口
talk is cheap show me the code
代码
限流配置RateLimitProperties
@Configuration
@ConfigurationProperties(prefix = "rateLimit")
@Data
public class RateLimitProperties {
private boolean enabled;
private List 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 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 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