一、目标
搞清楚Soul网关限流插件的运行原理,解读核心业务逻辑代码;
二、内容
2.1 背景
上一节我们一起学习了Soul网关限流插件的使用,这节我们就一起来看一下它背后运行的核心原理及关键代码实现;
Soul网关限流插件使用可以参考:
2.2 rateLimiter插件源码解析
2.2.1 RateLimiterPluginConfiguration解析
在Soul-admin开启了RateLimiter插件,并且配置了相关规则,在soul-bootstrap启动之后会自动加载配置类RateLimiterPluginConfiguration,自动向容器中注入限流插件RateLimiterPlugin;
@Configuration
public class RateLimiterPluginConfiguration {
@Bean
public SoulPlugin rateLimiterPlugin() {
return new RateLimiterPlugin(new RedisRateLimiter());
}
@Bean
public PluginDataHandler rateLimiterPluginDataHandler() {
return new RateLimiterPluginDataHandler();
}
}
2.2.2 RateLimiterPlugin解析
RateLimiterPlugin同样继承了AbstractSoulPlugin抽象类,关于AbstractSoulPlugin抽象类接口的分析,可以参考之前divide插件的分析文章,这里不再赘述。
Divide插件分析:https://blog.csdn.net/qq_38314459/article/details/112760726
这里重点看一下RateLimiterPlugin类里面的doExecute方法:
@Override
protected Mono<Void> doExecute(final ServerWebExchange exchange, final SoulPluginChain chain, final SelectorData selector, final RuleData rule) {
final String handle = rule.getHandle();
//取得配置参数
final RateLimiterHandle limiterHandle = GsonUtils.getInstance().fromJson(handle, RateLimiterHandle.class);
//根据 response.isAllowed() 来判断插件链是否继续执行
return redisRateLimiter.isAllowed(rule.getId(), limiterHandle.getReplenishRate(), limiterHandle.getBurstCapacity())
.flatMap(response -> {
if (!response.isAllowed()) {
exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
Object error = SoulResultWrap.error(SoulResultEnum.TOO_MANY_REQUESTS.getCode(), SoulResultEnum.TOO_MANY_REQUESTS.getMsg(), null);
return WebFluxResultUtils.result(exchange, error);
}
return chain.execute(exchange);
});
}
根据 response.isAllowed() 来判断插件链是否继续执行,如果是false直接抛出异常信息;
2.2.3 RedisRateLimiter解析
- 根据RuleID生成Keys
private static List<String> getKeys(final String id) {
String prefix = "request_rate_limiter.{" + id;
String tokenKey = prefix + "}.tokens";
String timestampKey = prefix + "}.timestamp";
return Arrays.asList(tokenKey, timestampKey);
}
- 组装读取lua脚本,脚本路径:/META-INF/scripts/request_rate_limiter.lua
private RedisScript<List<Long>> redisScript() {
DefaultRedisScript redisScript = new DefaultRedisScript<>();
redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("/META-INF/scripts/request_rate_limiter.lua")));
redisScript.setResultType(List.class);
return redisScript;
}
- 看懂上面两个方法,然后再看RedisRateLimiter.isAllowed()方法:
public Mono<RateLimiterResponse> isAllowed(final String id, final double replenishRate, final double burstCapacity) {
if (!this.initialized.get()) {
throw new IllegalStateException("RedisRateLimiter is not initialized");
}
//根据RuleID生成Keys
List<String> keys = getKeys(id);
//将keys 和 scriptArgs(速率,容量,当前时间戳(秒),当前需要的令牌数量)作为入参传给lua脚本
List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
//调用ReactiveRedisTemplate.execute()方法执行lua脚本
Flux<List<Long>> resultFlux = Singleton.INST.get(ReactiveRedisTemplate.class).execute(this.script, keys, scriptArgs);
return resultFlux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
.reduce(new ArrayList<Long>(), (longs, l) -> {
longs.addAll(l);
return longs;
}).map(results -> {
boolean allowed = results.get(0) == 1L;
Long tokensLeft = results.get(1);
RateLimiterResponse rateLimiterResponse = new RateLimiterResponse(allowed, tokensLeft);
log.info("RateLimiter response:{}", rateLimiterResponse.toString());
return rateLimiterResponse;
}).doOnError(throwable -> log.error("Error determining if user allowed from redis:{}", throwable.getMessage()));
}
- redis 执行lua脚本判断当前令牌桶剩余数量,并刷新令牌桶,返回:是否可以继续访问,令牌桶剩余容量;
2.2.4 request_rate_limiter.lua脚本
最后来一起看一下lua脚本文件
local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)
-- 参数1:速率
local rate = tonumber(ARGV[1])
-- 参数2:容量
local capacity = tonumber(ARGV[2])
--参数3:当前时间戳
local now = tonumber(ARGV[3])
--参数4:当前需要的令牌数量
local requested = tonumber(ARGV[4])
--填充时间=容量除以/速率
local fill_time = capacity/rate
--keys过期时间
local ttl = math.floor(fill_time*2)
-- 如果令牌为空,则填充当前容量
local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
last_tokens = capacity
end
--获取当前令牌时间戳,如果为nil,则设置为0
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
last_refreshed = 0
end
-- 当前时间戳和令牌最后刷新时间差值,和0比较,取最大值
local delta = math.max(0, now-last_refreshed)
--最后的令牌桶数量+当前填入数量(时间差*请求速率),和capacity比较取最小值,就是比较是不是快满了
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
--如果当前令牌桶没有满,则将令牌桶数量减1,说明当前可以继续请求,不丢弃
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
local allowed_num = 0
if allowed then
new_tokens = filled_tokens - requested
allowed_num = 1
end
-- 更新令牌桶当前容量和最新时间
redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)
return { allowed_num, new_tokens }
三、总结
rateLimiter 插件主要通过redis执行lua脚本来实现,保证原子操作。基本运行原理搞清楚了,如果要彻底理解,还要下来学习一下lua脚本相关知识。