跳到主要内容

设计限流器

问题​

如何用 Rust 实现一个高性能的限流器？

答案​

令牌桶算法​

use std::sync::Mutex;
use std::time::Instant;

pub struct TokenBucket {
    inner: Mutex<TokenBucketInner>,
}

struct TokenBucketInner {
    capacity: f64,       // 桶容量
    tokens: f64,         // 当前令牌数
    rate: f64,           // 每秒补充令牌数
    last_refill: Instant,
}

impl TokenBucket {
    pub fn new(capacity: f64, rate: f64) -> Self {
        Self {
            inner: Mutex::new(TokenBucketInner {
                capacity,
                tokens: capacity,
                rate,
                last_refill: Instant::now(),
            }),
        }
    }

    pub fn try_acquire(&self, tokens: f64) -> bool {
        let mut inner = self.inner.lock().unwrap();

        // 补充令牌
        let now = Instant::now();
        let elapsed = now.duration_since(inner.last_refill).as_secs_f64();
        inner.tokens = (inner.tokens + elapsed * inner.rate).min(inner.capacity);
        inner.last_refill = now;

        // 尝试消费
        if inner.tokens >= tokens {
            inner.tokens -= tokens;
            true
        } else {
            false
        }
    }
}

滑动窗口算法​

use std::collections::VecDeque;
use std::sync::Mutex;
use std::time::{Duration, Instant};

pub struct SlidingWindow {
    inner: Mutex<SlidingWindowInner>,
}

struct SlidingWindowInner {
    window: Duration,    // 窗口大小
    limit: usize,        // 窗口内最大请求数
    timestamps: VecDeque<Instant>,
}

impl SlidingWindow {
    pub fn new(window: Duration, limit: usize) -> Self {
        Self {
            inner: Mutex::new(SlidingWindowInner {
                window,
                limit,
                timestamps: VecDeque::new(),
            }),
        }
    }

    pub fn try_acquire(&self) -> bool {
        let mut inner = self.inner.lock().unwrap();
        let now = Instant::now();
        let cutoff = now - inner.window;

        // 清除过期的时间戳
        while inner.timestamps.front().map_or(false, |&t| t < cutoff) {
            inner.timestamps.pop_front();
        }

        if inner.timestamps.len() < inner.limit {
            inner.timestamps.push_back(now);
            true
        } else {
            false
        }
    }
}

算法对比​

算法优点缺点适用场景
令牌桶允许突发流量实现稍复杂API 限流
漏桶流量平滑不允许突发固定速率处理
滑动窗口精确统计内存随请求数增长小窗口限流
固定窗口实现简单边界突发粗粒度限流

Axum 中间件实现​

use axum::{middleware, extract::State, http::StatusCode, response::IntoResponse};
use std::sync::Arc;

async fn rate_limit_middleware(
    State(limiter): State<Arc<TokenBucket>>,
    request: axum::extract::Request,
    next: middleware::Next,
) -> impl IntoResponse {
    if limiter.try_acquire(1.0) {
        next.run(request).await.into_response()
    } else {
        StatusCode::TOO_MANY_REQUESTS.into_response()
    }
}

常见面试问题​

Q1: 分布式限流怎么做？​

答案：

单机限流用本地数据结构，分布式限流需要中心化存储：

  • Redis + Lua 脚本：原子操作，最常用
  • 滑动窗口 + Redis ZSET：精确但内存开销大
  • 本地 + 同步：每个节点本地限流 + 定期同步到 Redis

Rust 中可以用 redis crate 实现分布式令牌桶。

相关链接​