Soul网关源码分析-限流插件RateLimiter源码解读

一、简介

上一节我们一起学习了Soul网关限流插件的使用，这节我们就一起来看一下它背后运行的核心原理及关键代码实现；

二、源码分析

1.RateLimiterPluginConfiguration解析
在Soul-admin开启了RateLimiter插件，并且配置了相关规则，在soul-bootstrap启动之后会自动加载配置类RateLimiterPluginConfiguration，自动向容器中注入限流插件RateLimiterPlugin，如下所示：

@Configuration
public class RateLimiterPluginConfiguration {
    
    @Bean
    public SoulPlugin rateLimiterPlugin() {
        return new RateLimiterPlugin(new RedisRateLimiter());
    }
    
    @Bean
    public PluginDataHandler rateLimiterPluginDataHandler() {
        return new RateLimiterPluginDataHandler();
    }
}

2.RateLimiterPlugin解析
RateLimiterPlugin同样继承了AbstractSoulPlugin抽象类,重点看一下RateLimiterPlugin类里面的doExecute方法,如下：

@Override
    protected Mono<Void> doExecute(final ServerWebExchange exchange, final SoulPluginChain chain, final SelectorData selector, final RuleData rule) {
        final String handle = rule.getHandle();
        //取得配置参数
        final RateLimiterHandle limiterHandle = GsonUtils.getInstance().fromJson(handle, RateLimiterHandle.class);
        //根据 response.isAllowed() 来判断插件链是否继续执行
        return redisRateLimiter.isAllowed(rule.getId(), limiterHandle.getReplenishRate(), limiterHandle.getBurstCapacity())
                .flatMap(response -> {
                    if (!response.isAllowed()) {
                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                        Object error = SoulResultWrap.error(SoulResultEnum.TOO_MANY_REQUESTS.getCode(), SoulResultEnum.TOO_MANY_REQUESTS.getMsg(), null);
                        return WebFluxResultUtils.result(exchange, error);
                    }
                    return chain.execute(exchange);
                });
    }

根据 response.isAllowed() 来判断插件链是否继续执行，如果是false直接抛出异常信息。

3.RedisRateLimiter解析

• 根据RuleID生成Keys

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脚本，脚本路径：/META-INF/scripts/request_rate_limiter.lua

private RedisScript<List<Long>> redisScript() {
        DefaultRedisScript redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("/META-INF/scripts/request_rate_limiter.lua")));
        redisScript.setResultType(List.class);
        return redisScript;
    }

• 看懂上面两个方法，然后再看RedisRateLimiter.isAllowed()方法：

public Mono<RateLimiterResponse> isAllowed(final String id, final double replenishRate, final double burstCapacity) {
        if (!this.initialized.get()) {
            throw new IllegalStateException("RedisRateLimiter is not initialized");
        }
    //根据RuleID生成Keys
        List<String> keys = getKeys(id);
    //将keys 和 scriptArgs（速率，容量，当前时间戳（秒），当前需要的令牌数量）作为入参传给lua脚本
        List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", Instant.now().getEpochSecond() + "", "1");
    //调用ReactiveRedisTemplate.execute()方法执行lua脚本
        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()));
    }

redis 执行lua脚本判断当前令牌桶剩余数量，并刷新令牌桶，返回：是否可以继续访问，令牌桶剩余容量。

4.request_rate_limiter.lua脚本
最后来看一下lua脚本文件，如下：

local tokens_key = KEYS[1]
local timestamp_key = KEYS[2]
--redis.log(redis.LOG_WARNING, "tokens_key " .. tokens_key)
-- 参数1：速率
local rate = tonumber(ARGV[1])
-- 参数2：容量
local capacity = tonumber(ARGV[2])
--参数3：当前时间戳
local now = tonumber(ARGV[3])
--参数4：当前需要的令牌数量
local requested = tonumber(ARGV[4])
--填充时间=容量除以/速率
local fill_time = capacity/rate
--keys过期时间
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

--获取当前令牌时间戳，如果为nil，则设置为0
local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
  last_refreshed = 0
end

-- 当前时间戳和令牌最后刷新时间差值，和0比较，取最大值
local delta = math.max(0, now-last_refreshed)
--最后的令牌桶数量+当前填入数量（时间差*请求速率），和capacity比较取最小值，就是比较是不是快满了
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
--如果当前令牌桶没有满，则将令牌桶数量减1，说明当前可以继续请求，不丢弃
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.call("setex", tokens_key, ttl, new_tokens)
redis.call("setex", timestamp_key, ttl, now)

return { allowed_num, new_tokens }

总结

rateLimiter 插件主要通过redis执行lua脚本来实现，保证原子操作。基本运行原理搞清楚了，如果要彻底理解，还要下来学习一下lua脚本相关知识。