3、令牌桶算法实现：基于Python的令牌桶实现、参数调优、在Web服务中的应用

令牌桶算法，说白了就是给每个请求发一张「入场券」。我当年刚接触限流时，总觉得这玩意儿不就是个计数器吗？后来线上被突发流量教做人，才明白令牌桶的平滑能力有多重要。

这一章，咱们手写一个生产级的令牌桶，再聊聊参数怎么调，最后把它塞进Web服务里。你跟着我走一遍，以后面试、排障都能心里有底。

3.1 令牌桶的核心逻辑

先画个图，把脑子里的概念理清楚。令牌桶的核心就三个东西：桶、令牌、漏嘴。

核心公式：

  • 令牌生成速率 = r（个/秒）
  • 桶容量 = b（最大令牌数）
  • 当前令牌数 = min(上次令牌数 + 时间差 × r, b)
  • 请求通过条件：当前令牌数 ≥ 1
令牌桶算法核心流程 令牌生成器 速率 r 个/秒 持续注入 令牌桶 容量 b 当前令牌数 min(上次 + Δt×r, b) 请求消耗 请求通过 令牌 ≥ 1 请求拒绝 令牌不足 💡 关键特性： • 允许突发：桶满时最多可一次性消耗 b 个令牌，应对瞬时流量尖峰 • 长期平滑：平均速率被 r 限制，不会持续超载

你看这个图，令牌生成器像水龙头一样往桶里滴水，请求来了就从桶里舀一勺。桶满了水就溢出去，桶空了请求就等着或被拒绝。这就是令牌桶最朴素的形象。

3.2 手写一个生产级令牌桶

网上很多教程喜欢用定时器往桶里加令牌，我个人习惯不用定时器——用时间差计算更轻量，也更好控制。你想想看，定时器在高并发下反而容易成为瓶颈。

import time
import threading

class TokenBucket:
    """令牌桶限流器"""

    def __init__(self, rate: float, capacity: int):
        """
        :param rate: 令牌生成速率（个/秒）
        :param capacity: 桶容量（最大令牌数）
        """
        if rate <= 0 or capacity <= 0:
            raise ValueError("rate 和 capacity 必须大于 0")

        self.rate = rate
        self.capacity = capacity
        self._tokens = capacity          # 初始满桶
        self._last_time = time.monotonic()
        self._lock = threading.Lock()

    def _refill(self):
        """计算并补充令牌"""
        now = time.monotonic()
        elapsed = now - self._last_time
        # 核心公式：新增令牌 = 时间差 × 速率
        new_tokens = elapsed * self.rate
        self._tokens = min(self.capacity, self._tokens + new_tokens)
        self._last_time = now

    def acquire(self, tokens: int = 1) -> bool:
        """
        尝试获取 tokens 个令牌
        :return: True 表示获取成功，False 表示获取失败
        """
        if tokens <= 0:
            return True

        with self._lock:
            self._refill()
            if self._tokens >= tokens:
                self._tokens -= tokens
                return True
            return False

    def get_available_tokens(self) -> float:
        """获取当前可用令牌数（仅用于监控）"""
        with self._lock:
            self._refill()
            return self._tokens

我的经验：用 time.monotonic() 而不是 time.time()。系统时间可能被 NTP 回拨，monotonic 保证只增不减。我曾经因为这个坑，线上限流突然失效了半小时……

3.3 参数调优：rate 和 capacity 怎么配？

很多同学拿到令牌桶，第一个问题就是：rate 设多少？capacity 设多少？

说实话，没有万能公式。但我在几个项目里总结了一套经验法则，你可以参考。

场景 rate（QPS） capacity 说明
API 网关 业务峰值 × 1.2 rate × 2 ~ rate × 5 允许短时突发，但长期被 rate 限制
数据库连接池 连接池大小 连接池大小 × 1.5 防止瞬间建连打爆数据库
消息队列消费 消费者处理能力 rate × 10 容忍消费积压，但控制平均速率
登录接口 正常用户频率 × 3 rate × 1 严格限流，不允许突发

⚠️ 避坑指南：capacity 不要设得太大。我曾经给一个接口设了 capacity = rate × 100，结果凌晨定时任务一来，直接把后端打挂了。capacity 越大，突发流量越猛，后端扛不住。

3.4 在 Web 服务中实战

理论说完了，咱们把它塞进 Flask 里。嗯，这里要注意：中间件模式是最干净的，不要在每个路由里手动调 acquire。

from flask import Flask, request, jsonify
import time

app = Flask(__name__)

# 创建一个令牌桶：每秒 10 个请求，桶容量 20
bucket = TokenBucket(rate=10, capacity=20)

@app.before_request
def rate_limit():
    """全局限流中间件"""
    # 跳过健康检查等内部接口
    if request.path.startswith('/health'):
        return None

    if not bucket.acquire(1):
        return jsonify({
            "code": 429,
            "message": "Too Many Requests",
            "retry_after": 1.0 / bucket.rate
        }), 429, {"Retry-After": str(1.0 / bucket.rate)}

@app.route('/api/order', methods=['POST'])
def create_order():
    # 业务逻辑
    return jsonify({"status": "ok"})

@app.route('/api/search', methods=['GET'])
def search():
    # 业务逻辑
    return jsonify({"results": []})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

关键设计点：

  • 全局中间件 vs 路由装饰器：我建议用 before_request 做全局兜底，再在敏感路由上用装饰器做二次限流。分层防御更安全。
  • Retry-After 头：返回给客户端一个建议等待时间，客户端可以据此做退避。我见过太多客户端疯狂重试把服务打死的案例。
  • 健康检查豁免：负载均衡器的健康检查不能被限流，否则会被摘掉节点。

3.5 多维度限流：不只是全局

实际项目中，一个全局令牌桶往往不够。比如：每个用户、每个 IP、每个 API 都需要独立的限流策略。

我习惯用字典维护多个桶，配合 LRU 淘汰策略防止内存泄漏。

from collections import OrderedDict

class MultiBucketLimiter:
    """多维度令牌桶管理器"""

    def __init__(self, default_rate: float, default_capacity: int, max_buckets: int = 10000):
        self.default_rate = default_rate
        self.default_capacity = default_capacity
        self._buckets = OrderedDict()
        self._max_buckets = max_buckets
        self._lock = threading.Lock()

    def _get_bucket(self, key: str) -> TokenBucket:
        with self._lock:
            if key not in self._buckets:
                # 超过最大桶数，淘汰最久未使用的
                if len(self._buckets) >= self._max_buckets:
                    self._buckets.popitem(last=False)
                self._buckets[key] = TokenBucket(self.default_rate, self.default_capacity)
            else:
                # 移到末尾表示最近使用
                self._buckets.move_to_end(key)
            return self._buckets[key]

    def acquire(self, key: str, tokens: int = 1) -> bool:
        bucket = self._get_bucket(key)
        return bucket.acquire(tokens)

# 使用示例
limiter = MultiBucketLimiter(default_rate=5, default_capacity=10)

@app.before_request
def per_user_rate_limit():
    user_id = request.headers.get('X-User-Id', 'anonymous')
    if not limiter.acquire(user_id):
        return jsonify({"code": 429, "message": "用户级别限流"}), 429

我的经验：多维度限流一定要设置 max_buckets。我曾经没设上限，结果被爬虫遍历用户 ID，内存直接 OOM。LRU 淘汰 + 监控告警，这两样缺一不可。

3.6 性能与监控

令牌桶的性能瓶颈主要在锁竞争。单机场景下，Python 的 threading.Lock 够用。如果 QPS 超过 10 万，可以考虑用 Redis 实现分布式令牌桶，或者用 C 扩展。

监控方面，我建议至少暴露三个指标：

  • 当前令牌数：接近 0 说明系统快扛不住了
  • 拒绝请求数：突增说明需要扩容或调参
  • 平均获取等待时间：如果客户端等太久，用户体验会下降
# 简单的 Prometheus 指标暴露
from prometheus_client import Counter, Gauge

rejected_requests = Counter('rate_limit_rejected_total', '被限流拒绝的请求数')
available_tokens = Gauge('rate_limit_available_tokens', '当前可用令牌数')

@app.before_request
def monitor_rate_limit():
    if not bucket.acquire(1):
        rejected_requests.inc()
        return jsonify({"code": 429}), 429
    available_tokens.set(bucket.get_available_tokens())

说白了，令牌桶算法不难，难的是参数怎么调、怎么跟业务结合。你把这个实现吃透了，后面讲漏桶、滑动窗口、分布式限流，都是在这个基础上做变种。

← 上一章 📖 返回目录 下一章 →