65 lines
2.4 KiB
Java
65 lines
2.4 KiB
Java
package com.sf.vertx.handle;
|
|
|
|
import java.util.concurrent.TimeUnit;
|
|
|
|
import org.springframework.data.redis.core.RedisTemplate;
|
|
|
|
import com.sf.vertx.api.pojo.ApiCurrentLimitingConfig;
|
|
import com.sf.vertx.api.pojo.AppCurrentLimitingConfig;
|
|
import com.sf.vertx.constans.RedisKeyConfig;
|
|
import com.sf.vertx.service.impl.AppConfigServiceImpl;
|
|
import com.sf.vertx.utils.SpringUtils;
|
|
|
|
import cn.hutool.core.thread.ThreadUtil;
|
|
import io.vertx.ext.web.RoutingContext;
|
|
import lombok.extern.slf4j.Slf4j;
|
|
|
|
@Slf4j
|
|
public class RedisRateLimiter {
|
|
public Boolean acquire(RoutingContext rc, int pattern) {
|
|
String appCode = rc.request().headers().get(AppConfigServiceImpl.getSacAppHeaderKey());
|
|
// TODO 先测试app模式,后面通过app缓存获取模式
|
|
String key = null;
|
|
switch (pattern) {
|
|
case 1:
|
|
AppCurrentLimitingConfig appCurrentLimitingConfig = AppConfigServiceImpl
|
|
.getAppCurrentLimitingConfig(appCode);
|
|
key = RedisKeyConfig.APP_CURRENT_LIMITING_CONFIG_KEY + ":" + appCode;
|
|
return rateLimiter(key, appCode, appCurrentLimitingConfig.getThreshold(),
|
|
appCurrentLimitingConfig.getTimeWindow());
|
|
case 2:
|
|
ApiCurrentLimitingConfig apiCurrentLimitingConfig = AppConfigServiceImpl
|
|
.getApiCurrentLimitingConfig(appCode);
|
|
key = RedisKeyConfig.APP_CURRENT_LIMITING_CONFIG_KEY + ":" + appCode + ":" + rc.request().uri() + ":"
|
|
+ rc.request().method();
|
|
return rateLimiter(key, appCode, apiCurrentLimitingConfig.getThreshold(),
|
|
apiCurrentLimitingConfig.getTimeWindow());
|
|
default:
|
|
break;
|
|
}
|
|
return false;
|
|
|
|
}
|
|
|
|
@SuppressWarnings("unchecked")
|
|
private Boolean rateLimiter(String key, String appCode, Integer threshold, Integer timeWindow) {
|
|
RedisTemplate<String, Integer> redisTemplate = SpringUtils.getBean("redisTemplate", RedisTemplate.class);
|
|
Object count = redisTemplate.opsForValue().get(key);
|
|
// redisTemplate.delete(key);
|
|
log.info("count:{}", count);
|
|
// 初始化,设置过期时间
|
|
ThreadUtil.execAsync(() -> {
|
|
add(timeWindow, redisTemplate, key);
|
|
});
|
|
return count != null && Integer.valueOf(count.toString()) <= threshold ? true : false;
|
|
}
|
|
|
|
private void add(Integer timeWindow, RedisTemplate<String, Integer> redisTemplate, String key) {
|
|
// log.info("redis app threshold: {}", redisTemplate.opsForValue().get(key));
|
|
Long incr = redisTemplate.opsForValue().increment(key);
|
|
if (incr == 1) { // 创建,才设置时间窗口
|
|
redisTemplate.expire(key, timeWindow, TimeUnit.SECONDS);
|
|
}
|
|
}
|
|
}
|