Spring Cloud Gateway官方限流脚本分析

Spring Cloud Gateway是使用令牌桶算法来实现限流的，并采用Redis结合lua脚本的方式来实现分布式限流。

lua脚本地址：request_rate_limiter.lua

RedisRateLimiter：RedisRateLimiter.java

lua脚本如下：

--入参部分
local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)

local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

--填满空桶需要的时间
local fill_time = capacity/rate
--两倍的填满时间，作为key的失效时间。防止key太多，占用空间
local ttl = math.floor(fill_time*2)

--redis.log(redis.LOG_WARNING, "rate " .. ARGV[1])
--redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2])
--redis.log(redis.LOG_WARNING, "now " .. ARGV[3])
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

--当前桶中剩余令牌数。如果key不存在，初始化桶的容量数，即默认为满
local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
  last_tokens = capacity
end
--redis.log(redis.LOG_WARNING, "last_tokens " .. last_tokens)

--上次刷新时间。如果key不存在，初始化为0
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end
--redis.log(redis.LOG_WARNING, "last_refreshed " .. last_refreshed)

--距离上次限流的时间差
local delta = math.max(0, now-last_refreshed)
--时间差内能够恢复多少令牌
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
--申请的令牌足够
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens

--申请结果标识，0表示失败，1表示成功
local allowed_num = 0
--如果申请令牌通过，重新计算剩余令牌数。
if allowed then
  new_tokens = filled_tokens - requested
  allowed_num = 1
end

--redis.log(redis.LOG_WARNING, "delta " .. delta)
--redis.log(redis.LOG_WARNING, "filled_tokens " .. filled_tokens)
--redis.log(redis.LOG_WARNING, "allowed_num " .. allowed_num)
--redis.log(redis.LOG_WARNING, "new_tokens " .. new_tokens)

--将本次申请后剩余令牌数，和key更新时间，写回redis，并设置过期时间
redis.call("setex", tokens_key, ttl, new_tokens)
--将key更新时间，写回redis，并设置过期时间
redis.call("setex", timestamp_key, ttl, now)

--返回令牌申请结果，和剩余令牌数
return { allowed_num, new_tokens }

限流的逻辑还是比较简单的，脚本中我已经添加了注释，下面再稍加解释。

脚本开头是传入的参数：

• tokens_key：限流的标识。可以是ip，或者在spring cloud中可以是serviceID
• timestamp_key：令牌桶刷新的时间戳。
• rate ：令牌生产的速率
• capacity ：令牌桶的容量
• now ：当前时间戳
• requested：当前请求的令牌数，默认为1

下面是限流核心思想：

对于每个到来的请求，都需要获取令牌，默认申请一个令牌。如果桶中令牌数足够，就认为本次申请成功。申请成功后，需要记录下剩余的令牌数，和本次申请时间(戳)，这些信息会写回到redis。

需要注意两点：

一是这里剩余的令牌数的计算，除了减去本次申请消耗外，还要加上上次请求距今的时间内恢复的令牌数。令牌恢复的速度是传入的rate值，所以在上次请求与本次请求之间的时间差(delta)内，就会产生delta*rate个令牌。

二是记录的信息需要设置过期时间，主要是为了防止大量key占用空间。key过期后，下次再有请求到来时，会进行初始化，将桶中令牌数重置为满，将上次请求时间戳重置为0。这里的过期时间设置的是两倍的空桶装满令牌的需要耗费的时间，也就是2*(capacity/rate)。

 

再来看看调用lua脚本的代码：

/**
 * This uses a basic token bucket algorithm and relies on the fact that Redis scripts
 * execute atomically. No other operations can run between fetching the count and
 * writing the new count.
 */
@Override
@SuppressWarnings("unchecked")
public Mono<Response> isAllowed(String routeId, String id) {
    if (!this.initialized.get()) {
        throw new IllegalStateException("RedisRateLimiter is not initialized");
    }

    Config routeConfig = loadConfiguration(routeId);

    // How many requests per second do you want a user to be allowed to do?
    int replenishRate = routeConfig.getReplenishRate();

    // How much bursting do you want to allow?
    int burstCapacity = routeConfig.getBurstCapacity();

    try {
        List<String> keys = getKeys(id);

        // The arguments to the LUA script. time() returns unixtime in seconds.
        List<String> scriptArgs = Arrays.asList(replenishRate + "",
                burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
        // allowed, tokens_left = redis.eval(SCRIPT, keys, args)
        Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys,
                scriptArgs);
        // .log("redisratelimiter", Level.FINER);
        return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
                .reduce(new ArrayList<Long>(), (longs, l) -> {
                    longs.addAll(l);
                    return longs;
                }).map(results -> {
                    //【lua脚本中有两个返回】
                    //返回1表示表示获取令牌成功，返回0表示被限流
                    boolean allowed = results.get(0) == 1L;
                    //返回剩余的令牌数
                    Long tokensLeft = results.get(1);

                    Response response = new Response(allowed,
                            getHeaders(routeConfig, tokensLeft));

                    if (log.isDebugEnabled()) {
                        log.debug("response: " + response);
                    }
                    return response;
                });
    }
    catch (Exception e) {
        /*
         * We don't want a hard dependency on Redis to allow traffic. Make sure to set
         * an alert so you know if this is happening too much. Stripe's observed
         * failure rate is 0.01%.
         */
        log.error("Error determining if user allowed from redis", e);
    }
    return Mono.just(new Response(true, getHeaders(routeConfig, -1L)));
}            

传入给lua脚本的参数值是：

• replenishRate：与lua脚本中的入参rate对应
• burstCapacity：与lua脚本中的入参capacity对应
• Instant.now().getEpochSecond()：与lua脚本中的入参now对应
• 1：与lua脚本中的入参requested对应，默认为1

脚本的返回值决定了请求能够获得令牌，如果返回1表示获得令牌成功，如果返回0表示获取令牌失败，被限流。