|
@@ -1,14 +1,14 @@
|
|
|
-package com.diagbot.config.RateLimiter;
|
|
|
+package com.diagbot.config.redislimiter;
|
|
|
|
|
|
-import org.apache.commons.lang3.ObjectUtils;
|
|
|
import org.springframework.beans.BeansException;
|
|
|
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
|
|
|
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
|
|
|
+import org.springframework.cloud.gateway.support.ConfigurationService;
|
|
|
import org.springframework.context.ApplicationContext;
|
|
|
import org.springframework.context.ApplicationContextAware;
|
|
|
import org.springframework.data.redis.core.ReactiveRedisTemplate;
|
|
|
import org.springframework.data.redis.core.script.RedisScript;
|
|
|
-import org.springframework.validation.Validator;
|
|
|
+import org.springframework.util.ObjectUtils;
|
|
|
import org.springframework.validation.annotation.Validated;
|
|
|
import reactor.core.publisher.Flux;
|
|
|
import reactor.core.publisher.Mono;
|
|
@@ -22,13 +22,14 @@ import java.util.List;
|
|
|
import java.util.Map;
|
|
|
import java.util.concurrent.atomic.AtomicBoolean;
|
|
|
|
|
|
+
|
|
|
/**
|
|
|
* @Description:
|
|
|
* @author: gaodm
|
|
|
* @time: 2019/12/26 16:34
|
|
|
*/
|
|
|
public class SystemRedisRateLimiter extends AbstractRateLimiter<SystemRedisRateLimiter.Config> implements ApplicationContextAware {
|
|
|
- //这些变量全部从RedisRateLimiter复制的,都会用到。
|
|
|
+
|
|
|
public static final String REPLENISH_RATE_KEY = "replenishRate";
|
|
|
|
|
|
public static final String BURST_CAPACITY_KEY = "burstCapacity";
|
|
@@ -63,38 +64,37 @@ public class SystemRedisRateLimiter extends AbstractRateLimiter<SystemRedisRateL
|
|
|
private Config defaultConfig;
|
|
|
|
|
|
public SystemRedisRateLimiter(ReactiveRedisTemplate<String, String> redisTemplate,
|
|
|
- RedisScript<List<Long>> script, Validator validator) {
|
|
|
- super(Config.class, CONFIGURATION_PROPERTY_NAME, validator);
|
|
|
+ RedisScript<List<Long>> script, ConfigurationService configurationService) {
|
|
|
+ super(Config.class, CONFIGURATION_PROPERTY_NAME, configurationService);
|
|
|
this.redisTemplate = redisTemplate;
|
|
|
this.script = script;
|
|
|
initialized.compareAndSet(false, true);
|
|
|
}
|
|
|
|
|
|
- // public SystemRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity){
|
|
|
- // super(Config.class , CONFIGURATION_PROPERTY_NAME , null);
|
|
|
- // defaultConfig = new Config()
|
|
|
- // .setReplenishRate(defaultReplenishRate)
|
|
|
- // .setBurstCapacity(defaultBurstCapacity);
|
|
|
- //
|
|
|
- // }
|
|
|
- //具体限流实现,此处调用的是lua脚本
|
|
|
+ public SystemRedisRateLimiter(int defaultReplenishRate, int defaultBurstCapacity) {
|
|
|
+ super(Config.class, CONFIGURATION_PROPERTY_NAME, (ConfigurationService) null);
|
|
|
+ defaultConfig = new Config()
|
|
|
+ .setReplenishRate(defaultReplenishRate)
|
|
|
+ .setBurstCapacity(defaultBurstCapacity);
|
|
|
+
|
|
|
+ }
|
|
|
+
|
|
|
@Override
|
|
|
- public Mono<Response> isAllowed(String routeId, String id) {
|
|
|
+ public Mono<RateLimiter.Response> isAllowed(String routeId, String id) {
|
|
|
if (!this.initialized.get()) {
|
|
|
throw new IllegalStateException("RedisRateLimiter is not initialized");
|
|
|
}
|
|
|
if (ObjectUtils.isEmpty(rateLimiterConf)) {
|
|
|
throw new IllegalArgumentException("No Configuration found for route " + routeId);
|
|
|
}
|
|
|
- //获取的是自定义的map
|
|
|
Map<String, Integer> rateLimitMap = rateLimiterConf.getRateLimitMap();
|
|
|
//缓存的key
|
|
|
String replenishRateKey = routeId + "." + REPLENISH_RATE_KEY;
|
|
|
- //若map中不存在则采用默认值,存在则取值。
|
|
|
+
|
|
|
int replenishRate = ObjectUtils.isEmpty(rateLimitMap.get(replenishRateKey)) ? rateLimitMap.get(DEFAULT_REPLENISHRATE) : rateLimitMap.get(replenishRateKey);
|
|
|
//容量key
|
|
|
String burstCapacityKey = routeId + "." + BURST_CAPACITY_KEY;
|
|
|
- //若map中不存在则采用默认值,存在则取值。
|
|
|
+
|
|
|
int burstCapacity = ObjectUtils.isEmpty(rateLimitMap.get(burstCapacityKey)) ? rateLimitMap.get(DEFAULT_BURSTCAPACITY) : rateLimitMap.get(burstCapacityKey);
|
|
|
|
|
|
try {
|
|
@@ -104,7 +104,7 @@ public class SystemRedisRateLimiter extends AbstractRateLimiter<SystemRedisRateL
|
|
|
Instant.now().getEpochSecond() + "", "1");
|
|
|
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
|
|
|
|
|
|
- return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
|
|
|
+ return flux.onErrorResume((throwable) -> Flux.just(Arrays.asList(1L, -1L)))
|
|
|
.reduce(new ArrayList<Long>(), (longs, l) -> {
|
|
|
longs.addAll(l);
|
|
|
return longs;
|
|
@@ -143,7 +143,6 @@ public class SystemRedisRateLimiter extends AbstractRateLimiter<SystemRedisRateL
|
|
|
// this allows for using redis cluster
|
|
|
|
|
|
// Make a unique key per user.
|
|
|
- //此处可以自定义redis前缀信息
|
|
|
String prefix = "request_sys_rate_limiter.{" + id;
|
|
|
|
|
|
// You need two Redis keys for Token Bucket.
|