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 的高速率以几乎无损的延迟做到了限流的功能。