ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

用 golang.org/x/time/rate 给接口限流:令牌桶、突发容量与按 key 限速

用 golang.org/x/time/rate 给接口限流:令牌桶、突发容量与按 key 限速 用 golang.org/x/time/rate 给接口限流:令牌桶、突发容量与按 key 限速线上一个下游只扛得住 200 QPS,你的服务一激动把它打崩了;或者第三方 API 按调用次数收费,你想稳稳压在每秒 10 次。这种「平滑限速」不是用time.Sleep硬等能解决的——固定睡眠既浪费吞吐,又扛不住突发。Go 官方的golang.org/x/time/rate就是干这个的:一个久经验证的令牌桶实现。这篇把它的用法、突发容量、阻塞 vs 拒绝、以及最常见的「按用户 key 限流」一次讲透。先装依赖:go get golang.org/x/time/rate令牌桶两个参数:速率和桶容量rate.NewLimiter(r, b)只有两个参数,但很多人一上来就配错:r:令牌生成速率,类型是rate.Limit,单位是「每秒几个」。b:桶的容量(burst),即最多能攒多少个令牌,决定了瞬时能放过的最大突发量。packagemainimport(fmttimegolang.org/x/time/rate)funcmain(){// 每秒补 10 个令牌,桶最多攒 5 个 —— 稳态 10 QPS,允许一次冒 5 个的突发lim:rate.NewLimiter(rate.Limit(10),5)fori:0;i8;i{// Allow 立即返回:有令牌返 true 并扣一个,没有返 false 不阻塞ok:lim.Allow()fmt.Printf(req %d - %v\n,i,ok)}}运行会看到前 5 个是true(桶里攒的 5 个令牌一次放光),第 6 个开始基本是false——因为每 100ms(1/10 秒)才补一个令牌,循环跑得太快还没补上。这就是突发容量的意义:b控制「攒着的额度一次能爆发多少」,r控制「长期平均速率」。一个易错点:b至少要 1,否则一个都放不过去。想要「完全不允许突发、严格匀速」就把b设成 1。三种取令牌的姿势:Allow / Wait / ReserveAllow()只是最简单的一种。真实场景里你要根据「拿不到令牌时怎么办」来选:// 1) Allow:拿不到就立刻拒绝 —— 适合「宁可丢请求也不排队」的场景,比如防刷if!lim.Allow(){http.Error(w,rate limited,http.StatusTooManyRequests)return}// 2) Wait:拿不到就阻塞等,直到有令牌或 ctx 取消 —— 适合调用下游、爬虫这种「必须做完但要限速」funccallDownstream(ctx context.Context,lim*rate.Limiter)error{// 关键:Wait 一定要带 context,否则限速器抖动时会永久阻塞iferr:lim.Wait(ctx);err!nil{returnerr// ctx 超时/取消会在这返回}returndoRealCall()}// 3) Reserve:预约一个令牌,拿到需要等多久,自己决定等还是取消r:lim.Reserve()if!r.OK(){returnerrors.New(超出 burst 上限,永远等不到)}delay:r.Delay()// 还要等多久才能用ifdelay500*time.Millisecond{r.Cancel()// 太久就还回令牌,别占着returnerrTooBusy}time.Sleep(delay)Wait内部其实就是Reserve 按需 sleep,但它帮你处理了 ctx 取消,日常用它就够。记住Wait必须传context,这是最容易漏的:一旦上游请求已经超时,你还傻等令牌毫无意义,ctx 能让Wait立刻返回错误。一次批处理要 N 个令牌:AllowN / WaitN发一批消息、批量写下游,一次消耗多个额度时用带 N 的版本:// 这一批要发 3 条,一次扣 3 个令牌;不够就整批等/整批拒iferr:lim.WaitN(ctx,3);err!nil{returnerr}sendBatch(msgs)// 3 条注意N不能超过桶容量b,否则WaitN直接返回错误(桶根本装不下这么多,永远凑不齐)。批量大小和b要匹配着设。真实场景:按用户/IP 分别限流全局一个 limiter 会「一个用户刷爆,所有人一起被限」。正确做法是每个 key(用户 ID、IP)一个独立 limiter,用一个带过期清理的 map 管理:typekeyedLimiterstruct{mu sync.Mutex limitersmap[string]*entry r rate.Limit bint}typeentrystruct{lim*rate.Limiter seen time.Time// 最后使用时间,用于淘汰}funcnewKeyed(r rate.Limit,bint)*keyedLimiter{k:keyedLimiter{limiters:map[string]*entry{},r:r,b:b}gok.cleanupLoop()// 后台定期清理长期不活跃的 key,防止 map 无限膨胀returnk}func(k*keyedLimiter)get(keystring)*rate.Limiter{k.mu.Lock()deferk.mu.Unlock()e,ok:k.limiters[key]if!ok{eentry{lim:rate.NewLimiter(k.r,k.b)}k.limiters[key]e}e.seentime.Now()returne.lim}func(k*keyedLimiter)cleanupLoop(){forrangetime.Tick(time.Minute){k.mu.Lock()forkey,e:rangek.limiters{// 3 分钟没动的 key 直接删,回收内存iftime.Since(e.seen)3*time.Minute{delete(k.limiters,key)}}k.mu.Unlock()}}配成 HTTP 中间件:funcrateMiddleware(k*keyedLimiter,next http.Handler)http.Handler{returnhttp.HandlerFunc(func(w http.ResponseWriter,r*http.Request){ip:clientIP(r)// 从 X-Forwarded-For / RemoteAddr 取,注意别被伪造if!k.get(ip).Allow(){w.Header().Set(Retry-After,1)http.Error(w,too many requests,http.StatusTooManyRequests)return}next.ServeHTTP(w,r)})}别忘了那个后台清理 goroutine:如果每来一个新 IP 就建一个 limiter 却从不删除,map 会随攻击流量无限膨胀,最终 OOM。这是自己造轮子限流最容易埋的雷,x/time/rate只管单个桶,key 的生命周期得你自己兜。动态调速率:不用重启就改限流阈值限速阈值经常要按线上情况热调。SetLimit和SetBurst支持运行时改,配合配置中心很好用:// 下游反馈压力大,把速率从 10 降到 5,立即生效,不影响已在等的请求lim.SetLimit(rate.Limit(5))lim.SetBurst(3)还有个特殊值rate.Inf——表示不限流(无限速率),用来做「限流开关」比删中间件优雅:ifrateLimitDisabled{lim.SetLimit(rate.Inf)// 直接放行所有请求}小结rate.NewLimiter(r, b):r是每秒补多少个(长期均速),b是桶容量(突发上限),b至少为 1。拿不到令牌时:要拒绝用Allow,要排队用Wait(必须带 ctx),要自己权衡等多久用ReserveDelay。批量消耗用WaitN/AllowN,但N不能超过b。按用户/IP 限流 每 key 一个 limiter 后台过期清理,漏了清理就等着 OOM。热调阈值用SetLimit/SetBurst,rate.Inf当限流总开关。一句话记忆:令牌桶的r管长期均速、b管瞬时突发;按 key 限流别忘了给 map 配一把「扫地的扫帚」。
返回列表