限流算法的介绍见常见限流算法与使用redis实现简单的计数限流算法

通常来说，SpringCloud Gateway用于路由请求到微服务、过滤拦截请求或响应。但实际上，Gateway也能在路由层面上实现限流。具体的实现是Redis+lua实现令牌桶限流算法（调用lua的目的是实现原子操作），redis本身是支持调用lua脚本的。

原理

Gateway中的实现涉及下面两份源码：

filter/ratelimit/RedisRateLimiter.java

resources/META-INF/scripts/request_rate_limiter.lua

RedisRateLimiter.java的核心内容在isAllowed方法中，关键内容：

// 为不同的用户生成在redis中不同的键
List<String> keys = getKeys(id);
// 设置调用lua脚本的参数
List<String> scriptArgs = Arrays.asList(replenishRate + "",
        burstCapacity + "", Instant.now().getEpochSecond() + "",
        requestedTokens + "");
// 调用lua脚本
Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys,
        scriptArgs);

return flux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
    .reduce(new ArrayList<Long>(), (longs, l) -> {
        longs.addAll(l);
        return longs;
    }).map(results -> {
        // 如果返回的第一个结果是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;
    });

可以看出关键逻辑还是在lua脚本中的：

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]

-- 令牌产生速率
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
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

-- 上一次请求的时间戳
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end

-- 时间戳的差值
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
redis.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)

return { allowed_num, new_tokens }

其实lua脚本的逻辑还是挺清晰的，主要是通过redis存储了每个用户上一次请求剩余的令牌数和上一次请求的时间戳，并通过时间的差值计算令牌桶中的令牌数量。

存储两个内容的键分别是：

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

即存储令牌数的键为request_rate_limiter.{userId}.tokens，存储时间戳的键为request_rate_limiter.{userId}.timestamp。

使用

看完了原理我们就来看怎么使用，我就拿我已有的项目做实验。

引入依赖：

<!-- gateway-->
<dependency>
    <groupId>org.springframework.cloud</groupId>
    <artifactId>spring-cloud-starter-gateway</artifactId>
</dependency>
<!--redis-->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis-reactive</artifactId>
</dependency>

在网关微服务中进行配置：

配置一个KeyResolver的Bean对象，用于解析用户id，这里就把用户的IP地址作为id。

@Bean(name="ipKeyResolver")
public KeyResolver userKeyResolver() {
    return new KeyResolver() {
        @Override
        public Mono<String> resolve(ServerWebExchange exchange) {
            //获取远程客户端IP
            String hostName = exchange.getRequest().getRemoteAddress().getAddress().getHostAddress();
            System.out.println("hostName:"+hostName);
            return Mono.just(hostName);
        }
    };
}

在yaml文件中进行配置：

- id: service-edu
    uri: lb://service-edu
    predicates:
    - Path=/api/edu/**
    filters:
    - StripPrefix=1
    ## 这部分是限流的配置
    - name: RequestRateLimiter
      args:
        # 这个参数就是配置的KeyResolver的bean
        key-resolver: "#{@ipKeyResolver}"
        # 指的是令牌产生速度
        redis-rate-limiter.replenishRate: 1
        # 指的是令牌桶容量
        redis-rate-limiter.burstCapacity: 1

为了做验证我们把这两个参数都先设为1，下面进行测试：

可见限流成功，我们再在redis中查看是否有用户对应的键：

可见，键的名称也是和预期一致的。