Soul 网关源码分析(十二)ratelimiter 插件(二)

在上一篇我们熟悉了限流已经 ratelimiter 插件的实现原理,今天我们从源码入手详细地理一下整个流程。

源码分析

RedisRateLimiter类中:

public RedisRateLimiter() {
	//这里会调用 redisScript函数,读取并设置 this.script
    this.script = redisScript();
    initialized.compareAndSet(false, true);
}

/**
 *
 * @param id            规则ID
 * @param replenishRate 刷新率,就是每次填充进来的令牌数
 * @param burstCapacity 令牌桶的最大容量
 * @return {@code Mono<Response>} to indicate when request processing is complete
 */
@SuppressWarnings("unchecked")
public Mono<RateLimiterResponse> isAllowed(final String id, final double replenishRate, final double burstCapacity) {
	//如果没有初始化完成,则抛出异常
    if (!this.initialized.get()) {
        throw new IllegalStateException("RedisRateLimiter is not initialized");
    }
    List<String> keys = getKeys(id);
    //构造 redis 脚本参数,等会我们去lua脚本看看
    List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
    //通过 ReactiveRedisTemplate 来执行在构造函数中加载的脚本
    Flux<List<Long>> resultFlux = Singleton.INST.get(ReactiveRedisTemplate.class).execute(this.script, keys, scriptArgs);
    //当有异常发生时接收异常信息,输出和流中数据的类型相同的值,使用这个返回值替代异常的数据值返回给Subscriber
    return resultFlux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L))) 
    		//对返回的 List<Long> 做 reduce 操作
            .reduce(new ArrayList<Long>(), (longs, l) -> {
                longs.addAll(l);
                return longs;
            }).map(results -> { //针对流中的每一个元素,执行对应操作
                boolean allowed = results.get(0) == 1L; //如果 第一个结果是 1L,则允许请求通过
                Long tokensLeft = results.get(1); //获取桶中剩余的 token 数量
                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()));//在发生异常时打印日志
}

接下来我们看看 Lua 脚本中的内容:

--接收 tokens_key 和 timestamp_key
local tokens_key = KEYS[1] 
local timestamp_key = KEYS[2]
--将 rate, capacity, 当前秒级时间戳, 和 requested 参数传入,从Java代码中看到是 "1"
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])
-- 填充时间,计算出填充需要的时间,在上一节中我们的 capacity 设置为10,rate 为5 则填充时间为 2
local fill_time = capacity/rate
-- 对填充时间向上取整,获取 key 的过期时间
local ttl = math.floor(fill_time*2)

-- 获取剩余的 token 数量
local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
-- 如果不存在这个 key 则初始化令牌桶
  last_tokens = capacity
end
-- 获取最近更新的时间戳
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
-- 如果不存在这个key,则设置value 为 0
  last_refreshed = 0
end
-- 计算当前时间与上次更新时间的 时间差(秒)
local delta = math.max(0, now-last_refreshed)
-- 计算已经填充到桶中的令牌数,取桶容量 和 剩余令牌数加上时间差乘以填充速率 中**最小**的那个值
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
-- 计算是否允许请求和核心步骤:判断剩余令牌数是否大于等于1
local allowed = filled_tokens >= requested
-- 用一个变量存储更新后的 token数
local new_tokens = filled_tokens
local allowed_num = 0
-- 如果允许通过,则将减去 1 的令牌数更新到 tokens_key 中
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end
-- 调用 setex 来设置 key 过期时间 和 value
redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)
-- 返回表示状态的 allowed_num 和 剩余 token数
return { allowed_num, new_tokens }

下面是整个 限流计算的流程图:

总结

ratelimiter 插件巧妙地运用 Lua 脚本执行 redis 来实现,保证了操作的原子性,也借助 redis 的高速率以几乎无损的延迟做到了限流的功能。

相关推荐
©️2020 CSDN 皮肤主题: 大白 设计师:CSDN官方博客 返回首页