保护你的 API: Rust Axum 框架快速上手限流配置

我们在对接第三方接口时，常常会过于频繁地看到对方接口提示“超过调用次数限制”“超过调用频率限制”，或者直接返回 429 状态码。

在如今遍地都是 AI 服务的时代，这种情况更为明显。因为 AI 资源十分宝贵（主要是硬件显卡价格高昂导致并发成本非常高）。限制访问频率是充分利用 AI 资源最有效的手段。

漏桶算法（Leaky Bucket）

漏桶算法的基本思想是将请求比作水，当请求到达时，这些“水”会被存入一个固定容量的桶中。桶底有一个小孔，以恒定的速度向外漏水（处理请求）。如果流入桶中的水量超过了漏水的速度，多余的水就会溢出，即请求被拒绝1。这种方式可以有效地平滑流量，确保系统的输入流量不会超过设定的最大速率。

它的优点是什么？

1. 漏桶算法能够严格控制数据的传输速率，确保系统不会因突发流量而过载。
2. 算法逻辑简单，易于理解和实现。

它的缺点又是什么？

1. 由于其固定的流出速率，漏桶算法在面对突发流量时表现不佳，可能会导致大量请求被拒绝。
2. 即使网络中没有发生拥塞，漏桶算法也不能充分利用可用带宽，因为它总是按照预设的速率处理请求

令牌桶算法（Token Bucket）

令牌桶算法的核心思想是系统以恒定的速度向桶中添加令牌，每个请求需要消耗一个令牌才能被处理。如果桶中有足够的令牌，请求就可以立即得到处理；如果没有足够的令牌，请求将被拒绝或延迟。这种机制允许一定程度的突发流量，只要桶中有足够的令牌，请求就可以快速通过。

它的优点是什么？

1. 令牌桶算法可以在一定时间内处理超过平均速率的请求，提高了系统的灵活性和响应速度。
2. 更好地利用网络资源，因为它允许在空闲时段积累令牌，以便在高峰期使用。

它的缺点又是什么？

1. 相比于漏桶算法，复杂度高。
2. 如果桶中没有足够的令牌，请求可能会被延迟处理。

这里有一个对比表格

GCRA（Generic Cell Rate Algorithm )

一般直接使用漏桶算法的会比较少，因为它的缺点太明显了。GCRA算法是漏桶的生产改进版本，它有两种等价的描述方式：漏桶模式和虚拟调度模式。

漏桶模式

在这种模式下，GCRA定义了一个有限容量的桶，桶以每时间单位1的速度排水，并在每次确认请求后注水T。如果请求到达时桶内水位（X’）小于等于τ，则确认请求，否则拒绝请求。这种模式类似于传统的漏桶算法，但它引入了额外的参数来更好地处理突发流量。

虚拟调度模式

在这种模式中，GCRA通过比较请求的“理论到达时间（TAT, Theoretical Arrival Time）”与实际到达时间t来判定其能否被确认。算法定义了T，预期的请求间隔时间（可以推导出预期请求速度是1/T），以及容忍度τ，即请求可以早于TAT最多τ时间。如果t 小于 TAT-τ，则说明请求到达过早，不能被确认；请求被确认时，算法更新下一次的TAT=max(t,TAT)+T。这种模式更加优雅，便于深入理解和实现。

Axum 框架的使用

我们这里使用基于 GCRA 的限流中间件 tower_governor。tower_governor 是 tower 的一个中间件，用于实现接口限流。它可以基于对等 IP 地址、IP 地址标头、全局或通过自定义密钥的速率限制请求。 它是 governor 对 tower 的一个扩展。 governor 是 GCRA 的一个实现。

由于 Axum 是建立在 tower 层之上的，因此我们可以使用 tower_governor 来实现接口限流。

基于 IP 限流

假如我有一个组接口，，我想限制同一个IP内2秒内有5次请求配额。

// 基于 IP 限速，2秒内有5次请求配额
    let governor_conf = Arc::new(
        GovernorConfigBuilder::default()
            .per_second(2)
            .key_extractor(SmartIpKeyExtractor)
            .burst_size(5)
            .finish()
            .unwrap(),
    );
    let rate_limit_layer = GovernorLayer {
        config: governor_conf,
    };
    //限制登录接口组
    Router::new()
        .post("/auth", post(login).layer(rate_limit_layer))

tower_governor 提供了三个基本的提取器，

1. PeerIpKeyExtractor: 这是默认值，它使用请求的对等 IP 地址。
2. SmartIpKeyExtractor: 按顺序（x-forwarded-for、x-real-ip、forwarded）查找通常由反向代理提供的常见 IP 标识标头，并回退到对等 IP 地址。
3. GlobalKeyExtractor: 对所有传入请求使用相同的密钥。

基于用户登录凭证的限流

上面基于IP 限流只能是通用的场景。但是对于一些成本非常高的耗时场景，使用用户登录凭证限流才是更主要的做法。

对于基于用户登录凭证的限流，我们只要实现KeyExtractor接口就可以轻松实现一个自定义的提取器。

假设我们基于用户的 token 来限流，并且这个 token 是通过 header 头中的 Authorization: Bearer {token} 传递的。那我们们先定义一个结构体：

pub struct UserTokenExtractor;

实现KeyExtractor接口：

impl KeyExtractor for UserTokenExtractor {
    type Key = String;

    fn extract<B>(&self, req: &Request<B>) -> Result<Self::Key, GovernorError> {
        req.headers()
            .get("Authorization")
            .and_then(|token| token.to_str().ok())
            .and_then(|token| token.strip_prefix("Bearer ")).map(|token| token.trim().to_owned())
            .ok_or(GovernorError::Other {
                code: StatusCode::UNAUTHORIZED,
                msg: Some("You don't have permission to access".to_string()),
                headers: None,
            })
    }
    fn key_name(&self, key: &Self::Key) -> Option<String> {
        Some(key.to_string())
    }
    fn name(&self) -> &'static str {
        "UserToken"
    }
}

最主要的是 extract 的方法，用于从HTTP请求中提取授权令牌。具体功能如下：

1. 从请求头中获取 Authorization 字段。
2. 检查字段值是否为有效的字符串。
3. 去除字符串前缀 Bearer 。
4. 去除前后空白字符并返回。
5. 如果任何一步失败，返回 GovernorError 错误。

使用

// 基于用户 token 限速，5秒内有3次请求配额
    let token_rate_limit_layer = GovernorLayer {
        config: Arc::new(
            GovernorConfigBuilder::default()
            .per_second(5)
            .burst_size(3)
            .key_extractor(UserTokenExtractor)
            .use_headers()
            .finish()
            .unwrap(),
        ),
    };
    //限制生成接口的请求速率
    Router::new()
        .post("/generate", post(generate).layer(token_rate_limit_layer))

当我们发起请求过快时，会出现以下错误：

Too Many Requests! Wait for 1s