Soul网关源码阅读（十五） RateLimit插件使用

Soul网关源码阅读（十五） RateLimit插件使用

RateLimiter原理

RateLimiter是基于Spring Could RateLimter来实现的，RateLimiter是基于令牌桶的限流算法。所以理解限流核心是理解令牌桶算法，然后奖这个算法应用到网关的拦截器上。常用的限流算法有：漏桶算法和令牌桶算法。由于漏桶算法需要设置流出速率，接口响应速度，所以不能灵活反应当前系统压力，令牌桶算法是更好的选择。

漏桶算法

漏桶(Leaky Bucket)算法思路很简单,水(请求)先进入到漏桶里,漏桶以一定的速度出水(接口有响应速率),当水流入速度过大会直接溢出(访问频率超过接口响应速率),然后就拒绝请求,可以看出漏桶算法能强行限制数据的传输速率.示意图如下:

令牌桶算法

令牌桶算法(Token Bucket)和 Leaky Bucket 效果一样但方向相反的算法,更加容易理解.随着时间流逝,系统会按恒定1/QPS时间间隔(如果QPS=100,则间隔是10ms)往桶里加入Token(想象和漏洞漏水相反,有个水龙头在不断的加水),如果桶已经满了就不再加了.新请求来临时,会各自拿走一个Token,如果没有Token可拿了就阻塞或者拒绝服务.

所以需要注意两个算法的区别：

都是基于桶：

• 漏桶的桶是用来装请求（request）的，水龙头是用户的请求来源，它是不受控制的，当进来的请求个数多余桶容量，多余的请求将会被拒绝。
• 令牌桶的桶是用来装token的，水龙头是我们生成令牌的来源，它是可以调节的，只要桶里面还有令牌，既可以继续服务。
• 简单的说，漏桶入口不受限，出口受受限；令牌桶入口受限，出口不受限。

使用步骤

启动Ratelimiter插件

启动admin后台，在后台服务中插件管理打开ratelimit插件，需要配置redis连接信息，rateLimit使用的令牌桶基于redis。

为RateLimiter插件添加selector和rule

这里我们后端启动一个http服务作为测试。

配置一个匹配http下/order/**接口的规则，ratelimiter对这些匹配上的接口适用，为方便测试容量设置为5，rate设置为1。

启动soul网关添加对应配置

<!-- soul ratelimiter plugin start-->
  <dependency>
      <groupId>org.dromara</groupId>
      <artifactId>soul-spring-boot-starter-plugin-ratelimiter</artifactId>
      <version>${last.version}</version>
  </dependency>
  <!-- soul ratelimiter plugin end-->

接口测试

当请求速度超过设置的阈值过后，请求将会被拦截。

源码分析

先看看starter模块，配置了一些什么内容。

RateLimiterPluginConfiguration配置类

@Configuration
public class RateLimiterPluginConfiguration {
    
    /**
     * RateLimiter plugin.
     *
     * @return the soul plugin
     */
    @Bean
    public SoulPlugin rateLimiterPlugin() {
        return new RateLimiterPlugin(new RedisRateLimiter());
    }
    
    /**
     * Rate limiter plugin data handler plugin data handler.
     *
     * @return the plugin data handler
     */
    @Bean
    public PluginDataHandler rateLimiterPluginDataHandler() {
        return new RateLimiterPluginDataHandler();
    }
}

• RateLimiterPlugin

  RateLimiter插件，组合了RedisRateLimiter。

• PluginDataHandler

  数据处理器，用来处理admin同步过来的数据，他的主要工作是进行redis的配置。

看来重点在RedisRateLimiter这个类。

RedisRateLimiter

@Slf4j
public class RedisRateLimiter {
		//redis脚本
    private final RedisScript<List<Long>> script;
		//是否初始化
    private final AtomicBoolean initialized = new AtomicBoolean(false);

    /**
     * Instantiates a new Redis rate limiter.
     */
    public RedisRateLimiter() {
        this.script = redisScript(); //从classPath加载lua脚本
        initialized.compareAndSet(false, true);
    }

    /**
     * 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.
     * 这个方法用来从获取redis获取令牌，该执行方法依靠lua脚本进行原子操作，在对redis桶令牌进行获取和写入期间，其他操作将会阻塞。
     * @param id            is rule id
     * @param replenishRate replenishRate 生产速度1
     * @param burstCapacity burstCapacity 容量之前设置的5
     * @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); //获取tokenKey和timestampKey
        List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");//获取lua脚本需要的依赖参数，刷新速度，桶大小，当前时间，请求个数1
        Flux<List<Long>> resultFlux = Singleton.INST.get(ReactiveRedisTemplate.class).execute(this.script, keys, scriptArgs);
        return resultFlux.onErrorResume(throwable -> Flux.just(Arrays.asList(1L, -1L)))
                .reduce(new ArrayList<Long>(), (longs, l) -> {
                    longs.addAll(l);
                    return longs;
                }).map(results -> {
                    boolean allowed = results.get(0) == 1L;
                    Long tokensLeft = results.get(1);
                    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()));
    }

    private static List<String> getKeys(final String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }
    
	 //...
}

Lua脚本

上面isAllowed方法是整个RateLimiter限流的核心方法，其原理是通过redis来执行lua脚本判断是否可以获取令牌。

Lua脚本的参数分别为 ：

• 令牌桶的刷新速度
• 桶大小
• 以秒为单位的时间戳
• 桶流出速度（请求个数，即为1）

Lua脚本如下：

local tokens_key = KEYS[1] # token的key
local timestamp_key = KEYS[2] # 时间戳的key
--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 #桶容量除以生成速度,满桶所需时间（换言之也就是消耗完令牌所需时间）
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)

local last_tokens = tonumber(redis.call("get", tokens_key)) #获取token
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)) ##获取上一次刷新时间，如果为空为0
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) ## 上一次更新时间的增量，如果首次则为0
## 装入的token= 上次的token加上增量时间*速率。 比如上次还剩50个token（100的容量），速率为10*，那么3秒过后，应该添加是80
local filled_tokens = math.min(capacity, last_tokens+(delta*rate)) 
local allowed = filled_tokens >= requested ## 如果填入的token大于请求的个数，则为true
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)

redis.call("setex", tokens_key, ttl, new_tokens) # 使用setex 设置可获取token，及过期时间
redis.call("setex", timestamp_key, ttl, now) # 设置当前过期时间

return { allowed_num, new_tokens } ## 返回是否可用bool，以及剩下可用的token

总结

限流的关键的在于搞懂令牌桶算法的实现，如何利用redis的lua脚本来直线令牌桶算法，看一下他的实现流程，如下图

实现细节需要考虑：

• 如何实现以恒定速率往桶中放令牌？

  定时任务去更新行不行？可以，但没必要，因为这样增加了实现难度添加了系统负担。

  更好的解决办法就是，使用他的时候再去计算delta时间内应该增加的token值，并更新可用token。

• 为何要同时维护两个key，一个token_key一个timestamp_key？

  token_key用来保存上一次请求过后桶中剩下的token，timestamp_key用来计算本次应该新增的token数，需要用（当前时间-上一次更新数据）*速率。

• 为什么需要为token_key设置过期时间，且每次请求都更新ttl过期时间为capacity/rate*2？

  这个问题交给读者思考一下。