你的API每秒能扛住多少请求?这个数字很多开发者直到被刷量攻击时才真正关心。2023年Cloudflare统计,应用层DDoS攻击同比增长了67%,而大多数初创团队的API连基础的流量控制都没做。
限流(Rate Limiting)不是什么新鲜概念,但实现方式千差万别。固定窗口计数器简单粗暴,却在整点重置时容易被瞬间击穿;滑动窗口日志精确,内存开销却让人肉疼。Token Bucket(令牌桶)算法在这两者之间找到了一个微妙的平衡点——允许突发流量,同时守住长期吞吐的底线。
今天这篇不是调包教程。我们用纯Python手搓一个Token Bucket限流器,塞进FastAPI的中间件,再给它加上按用户区分的追踪能力。最后你会拿到一段可以直接复制到生产环境的代码。
Token Bucket的工作原理:一个会自己续杯的桶
想象一个容量固定的桶,桶里装着令牌。每个请求进来时,必须从桶里拿走一枚令牌才能通过。桶空了?请求直接吃429(Too Many Requests)。
关键设计在于"续杯"机制:无论有没有请求进来,系统都会以固定速率往桶里加令牌,直到装满为止。这意味着你可以瞬间爆发——比如桶容量10,你可以一口气打10个请求——但爆发后必须等令牌慢慢恢复,才能继续。
两个参数决定一切:capacity(桶容量)控制突发上限,refill_rate(每秒补充令牌数)决定长期吞吐量。
对比其他算法,Token Bucket的优势很明显。固定窗口在每分钟重置时,恶意用户可以在59秒和01秒各打一波,实际吞掉两倍配额;滑动窗口没这个问题,但需要记录每个请求的时间戳,内存随流量线性增长。Token Bucket只用两个变量(当前令牌数、上次补充时间),O(1)空间复杂度搞定一切。
手搓核心:一个不到50行的Python类
开始写代码。我们需要一个类来管理每个用户的令牌桶状态。核心逻辑只有三件事:初始化时设置容量和补充速率;每次请求时检查并扣减令牌;按时间自动补充。
关键点在于线程安全。FastAPI跑在异步环境里,多个请求可能同时操作同一个桶。这里用asyncio.Lock保证原子性——检查令牌和扣减必须一气呵成,否则会出现超卖。
代码结构很直白:__init__里存capacity、refill_rate、tokens(当前数量)、last_update(上次补充时间戳)、还有一把锁。consume方法先算距离上次更新过了多久,按refill_rate补充令牌(不超过capacity),然后判断够不够扣。
补充令牌的计算要注意浮点数精度。直接用time.time()的差值乘refill_rate,结果可能带小数。这里选择保留小数部分的令牌——比如0.5个令牌虽然不能让请求通过,但累积到1.0时就可以。这种"信用"机制让长期平均速率严格等于refill_rate,避免整数取整带来的系统性偏差。
塞进FastAPI:中间件实现按用户限流
单个桶不够,生产环境需要按用户区分。最简单的方案是用API Key或者用户ID做键,每个键对应一个TokenBucket实例。中间件里先识别用户,再找对应的桶,最后决定是否放行。
用户识别策略取决于你的认证方式。如果用了OAuth2,可以从Authorization头解析;如果是内部服务,X-API-Key更常见。这里演示一个通用方案:优先查自定义头, fallback到IP地址。注意IP不能直接用request.client.host——经过反向代理后这个值可能是127.0.0.1,要按X-Forwarded-For链取第一个非内网地址。
中间件的结构遵循ASGI规范。外层函数接收app,返回一个async def middleware(request, call_next)的协程。内部维护一个全局字典buckets,键是用户标识,值是TokenBucket实例。字典访问也要加锁,防止并发创建重复实例。
限流触发时,标准做法是返回429状态码,并在响应头里告知客户端限制详情。RFC 6585定义了几个标准头:X-RateLimit-Limit(总配额)、X-RateLimit-Remaining(剩余)、X-RateLimit-Reset(重置时间)。虽然这些头是"X-"前缀的自定义扩展,但GitHub、Twitter等主流平台都在用,已经形成事实标准。
我们的实现会在响应头里注入这三个值。Reset时间戳的计算稍 tricky:用当前令牌缺口除以refill_rate,向上取整到秒。客户端拿到后可以据此做退避重试。
测试与调参:别让限流变成误伤
代码写完需要验证。写个简单的脚本,用httpx或者aiohttp并发打请求,观察429出现的时机是否符合预期。
调参是门手艺。capacity设太小,正常用户的批量操作会被误杀;设太大,又失去保护作用。一个经验法则:capacity至少覆盖用户单次交互的最大请求数。比如你的页面加载需要5个API调用,那capacity不能低于5。refill_rate则根据业务峰值定,日常流量的2-3倍通常够用。
还要考虑冷启动问题。新用户的桶是满的,这意味着攻击者可以不断换身份来绕过限制。配合API Key的申请审核、或者IP级别的二级限流,才能堵住这个口子。
存储层也有优化空间。目前的实现把桶状态存在内存里,服务重启就清零,多实例部署时各管各的。要解决这个问题,需要把状态外置到Redis,用Lua脚本保证检查-扣减的原子性。那又是另一个故事了。
你的API现在能扛住多少QPS?如果还没测过,建议今晚就压一压。毕竟刷量攻击不会提前打招呼,而第一封用户投诉邮件往往比监控告警来得更早。
热门跟贴