如何动态调整golang.org/x/time/rate限流器配额并实现预留令牌归还

根本原因 #

golang.org/x/time/rate库的Reserve()返回的Reservation调用Cancel()时，只会回退还未到可用时间的预留令牌。如果res.Delay() == 0，说明你已经拿到了当前可直接消耗的令牌，库默认视为已经被消耗，此时调用Cancel()不会触发回退逻辑。

解决方案 #

可以利用库本身的预留取消逻辑，实现已消耗令牌的手动退还，不需要修改库源码，完全兼容官方API。

实现退还令牌的工具函数 #

import (
    "time"
    "golang.org/x/time/rate"
)

// returnRateTokens 向限流器退还n个已消耗的令牌，最多不超过限流器的突发上限
func returnRateTokens(lim *rate.Limiter, n int) {
    if n <= 0 {
        return
    }
    // 预约未来足够久的n个令牌，会优先扣除当前可用令牌
    res := lim.ReserveN(time.Now().Add(24 * time.Hour), n)
    if !res.OK() {
        return
    }
    // 取消预约，所有被扣除/预留的令牌都会被退回令牌池
    res.Cancel()
}

业务逻辑改造示例 #

原来用Reserve()判断即时可用的逻辑可以替换为更简单的Allow()，配合退还函数实现重试场景不占配额的需求：

lim := rate.NewLimiter(rateLimit, burstLimit)

for range time.NewTicker(time.Second).C {
    fmt.Println(time.Now(), "tick")
    go func() {
        // 无可用即时令牌直接丢弃请求
        if !lim.Allow() {
            fmt.Println(time.Now(), "dropped")
            return
        }

        // 执行业务请求
        err := callYourAPI()
        // 判断是否为幂等可重试的错误
        if err != nil && isRetryableError(err) {
            // 退还已消耗的令牌，重试请求不占配额
            returnRateTokens(lim, 1)
            // 此处可添加你的重试逻辑
            fmt.Println(time.Now(), "request retryable, token returned")
            return
        }

        fmt.Println(time.Now(), "request done")
    }()
}

注意事项 #

• 该方案完全基于官方库的公开API实现，没有依赖私有字段，不会因为库版本升级失效
• 退还的令牌数量会自动被限制在限流器的突发（burst）上限内，不会出现令牌溢出问题
• 如果你的场景需要单次预留/退还多个令牌，修改函数的n参数即可

内容的提问来源于stack exchange，提问作者andig