Spring Gateway限流源码解析

Spring Gateway限流

Spring Gateway的限流是通过过滤器来实现的，这个过滤器首先会从配置中拿到

Key解析器
限流器
是否拒绝空Key
空Key的状态码

java 复制代码

// Key解析器
KeyResolver resolver = getOrDefault(config.keyResolver, defaultKeyResolver);
// 限流器
RateLimiter<Object> limiter = getOrDefault(config.rateLimiter, defaultRateLimiter);
// 是否拒绝空Key
boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
// 拒绝空Key的状态码的Holder
HttpStatusHolder emptyKeyStatus = HttpStatusHolder
    .parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));

说白了就是下面这些配置的值：

yaml 复制代码

spring:
  cloud:
    gateway:
      routes:
      - id: requestratelimiter_route
        uri: https://example.org
        filters:
        - name: RequestRateLimiter
          args:
            # 补充令牌的速度
            redis-rate-limiter.replenishRate: 10
            # 令牌桶容量
            redis-rate-limiter.burstCapacity: 20
            # 一次请求消耗的令牌数
            redis-rate-limiter.requestedTokens: 1

当然，如果用户没有配置，那么就会拿到默认值。

ServerWebExchange

首先得了解ServerWebExchange是HTTP请求-响应交互的约定。

它提供对HTTP请求和响应的访问，还公开其他与服务器端处理相关的属性和功能，如请求属性。

也就是对一次请求-响应进行了封装。

Key解析器

KeyResolver用于解析请求该请求的Key，也就是这个请求的唯一标识，不过一个请求的唯一标识可以有很多，可以是他的ip，也可以是他的某个请求头比如token，所以这里将Key解析器抽象为一个接口：

java 复制代码

public interface KeyResolver {
	Mono<String> resolve(ServerWebExchange exchange);
}

返回Mono是因为在Spring WebFlux 中，响应通常是异步的，这意味着它们不会立即返回，而是需要一些时间来处理，而Gateway基于Spring WebFlux，所以自然也是异步的，从而需要返回一个Mono来获取结果。

Spring默认提供了一个PrincipalNameKeyResolver的实现类，它将每一次请求的Principal的name，也就是该请求已验证用户的name属性来充当key。

java 复制代码

public class PrincipalNameKeyResolver implements KeyResolver {
	public static final String BEAN_NAME = "principalNameKeyResolver";

	@Override
	public Mono<String> resolve(ServerWebExchange exchange) {
		return exchange.getPrincipal().flatMap(p -> Mono.justOrEmpty(p.getName()));
	}
}

限流器

限流器是整个Gateway实现限流的核心，它规定了一个用于判断改接口是否被允许的方法：

java 复制代码

public interface RateLimiter<C> extends StatefulConfigurable<C> {
    // 判断某一条路由规则的某一个请求是否被允许
    // routeId: 路由规则的唯一标识
    // id: 请求的唯一标识 -> 也就是用KeyResolver解析出来的Key
	Mono<Response> isAllowed(String routeId, String id);
	
    // 封装的返回值
	class Response {
        // 请求是否被允许
        private final boolean allowed;
        // 令牌的剩余量
		private final long tokensRemaining;
        // 响应的请求头
		private final Map<String, String> headers;
        // ...
	}
}

Spring提供了一个默认的实现类RedisRateLimiter，这个限流器是采用Redis + Lua脚本利用令牌桶的方式来实现限流的，首先它会从配置中拿到三个参数，分别对应着：

令牌填充速率
桶容量
每次请求消耗令牌数

java 复制代码

// 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();

// How many tokens are requested per request?
int requestedTokens = routeConfig.getRequestedTokens();

由于令牌桶的逻辑是在Redis中进行的，所以需要在Redis执行Lua脚本，而之所以要在Redis层面实现，是因为这样可以减少频繁地访问Redis从而因为频繁的网络请求导致性能损耗。

在执行Redis的Lua脚本之前，首先需要获取参数的Keys和Args列表：

java 复制代码

List<String> keys = getKeys(id);

// The arguments to the LUA script. time() returns unixtime in seconds.
List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", "", requestedTokens + "");

// allowed, tokens_left = redis.eval(SCRIPT, keys, args)
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);

其中getKeys()的具体实现如下：

java 复制代码

static List<String> getKeys(String id) {
    // use `{}` around keys to use Redis Key hash tags
    // this allows for using redis cluster

    // Make a unique key per user.
    String prefix = "request_rate_limiter.{" + id;

    // You need two Redis keys for Token Bucket.
    String tokenKey = prefix + "}.tokens";
    String timestampKey = prefix + "}.timestamp";
    return Arrays.asList(tokenKey, timestampKey);
}

不难看出最后能得到两个Redis的Key，假设ID为1：

request_rate_limiter.{1}.tokens
request_rate_limiter.{1}.timestamp

其中request_rate_limiter.{1}.tokens存储的是令牌的数量

而request_rate_limiter.{1}.timestamp则存的是上一次请求的时间戳

令牌桶算法

令牌桶算法是想象有一个桶，桶里会存放令牌，请求需要通过，首先得拿到指定数量的令牌才能通过，当然，桶里会源源不断地补充令牌。

假设桶的容量是capacity，而补充令牌的速度为rate，允许通过的令牌数为1个，那么可以写出如下伪代码：

ini 复制代码

capacity = 20
rate = 10
nowTime = now()
NEED_TOKEN_NUM = 1

lastTime = getLastTime()
// 获取剩余的令牌
remainingTokens = getRemainingTokens()
// 上次剩余的加上上次到这次之间补充的
tokens = remainingTokens + rate * (lastTime - nowTime)
// 不得超过最大容量
tokens = min(tokens, capacity)
if (tokens >= NEED_TOKEN_NUM) {
    // 取得需要的令牌
    tokens -= NEED_TOKEN_NUM
    // 设置上一次时间为当前时间
    setLastTime(nowTime)
    return true
} else {
    return false
}

Gateway的Lua脚本如下：

lua 复制代码

redis.replicate_commands()

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 = redis.call('TIME')[1]
local requested = tonumber(ARGV[4])

local fill_time = capacity/rate
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 " .. now)
--redis.log(redis.LOG_WARNING, "requested " .. ARGV[4])
--redis.log(redis.LOG_WARNING, "filltime " .. fill_time)
--redis.log(redis.LOG_WARNING, "ttl " .. ttl)

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)

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
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)

if ttl > 0 then
  redis.call("setex", tokens_key, ttl, new_tokens)
  redis.call("setex", timestamp_key, ttl, now)
end

-- return { allowed_num, new_tokens, capacity, filled_tokens, requested, new_tokens }
return { allowed_num, new_tokens }

其中值得注意的是TTL的设置是根据capacity/rate来确定的，说白了就是桶被填满的时间。

是否拒绝空Key

如果denyEmpty为true也就是配置了拒绝空Key则会在后面当解析出的Key为空的时候直接拒绝请求，否则不限流直接放行：

java 复制代码

if (EMPTY_KEY.equals(key)) {
if (denyEmpty) {
    setResponseStatus(exchange, emptyKeyStatus);
    return exchange.getResponse().setComplete();
}
	return chain.filter(exchange);
}