本篇文章承接上一篇[go-zero 如何扛住流量衝擊（一）](https://gocn.vip/topics/11247)。 上一篇介紹的是 `go-zero` 中滑動視窗限流，本篇介紹另外一個 `tokenlimit` ，令牌桶限流。 ## 使用 ```go const ( burst = 100 rate = 100 seconds = 5 ) store := redis.NewRedis("localhost:6379", "node", "") fmt.Println(store.Ping()) // New tokenLimiter limiter := limit.NewTokenLimiter(rate, burst, store, "rate-test") timer := time.NewTimer(time.Second * seconds) quit := make(chan struct{}) defer timer.Stop() go func() { <-timer.C close(quit) }() var allowed, denied int32 var wait sync.WaitGroup for i := 0; i < runtime.NumCPU(); i++ { wait.Add(1) go func() { for { select { case <-quit: wait.Done() return default: if limiter.Allow() { atomic.AddInt32(&allowed, 1) } else { atomic.AddInt32(&denied, 1) } } } }() } wait.Wait() fmt.Printf("allowed: %d, denied: %d, qps: %d\n", allowed, denied, (allowed+denied)/seconds) ``` ## tokenlimit 從整體上令牌桶生產token邏輯如下： - 使用者配置的平均傳送速率為r，則每隔1/r秒一個令牌被加入到桶中； - 假設桶中最多可以存放b個令牌。如果令牌到達時令牌桶已經滿了，那麼這個令牌會被丟棄； - 當流量以速率v進入，從桶中以速率v取令牌，拿到令牌的流量通過，拿不到令牌流量不通過，執行熔斷邏輯； `go-zero` 在兩類限流器下都採取 `lua script` 的方式，依賴redis可以做到分散式限流，`lua script`同時可以做到對 token 生產讀取操作的原子性。 下面來看看 `lua script` 控制的幾個關鍵屬性： | **argument** | **mean** | | ------------ | ------------------------------------------ | | ARGV[1] | rate 「每秒生成幾個令牌」 | | ARGV[2] | burst 「令牌桶最大值」 | | ARGV[3] | now_time「當前時間戳」 | | ARGV[4] | get token nums 「開發者需要獲取的token數」 | | KEYS[1] | 表示資源的tokenkey | | KEYS[2] | 表示重新整理時間的key | ```lua -- 返回是否可以活獲得預期的token local rate = tonumber(ARGV[1]) local capacity = tonumber(ARGV[2]) local now = tonumber(ARGV[3]) local requested = tonumber(ARGV[4]) -- fill_time：需要填滿 token_bucket 需要多久 local fill_time = capacity/rate -- 將填充時間向下取整 local ttl = math.floor(fill_time*2) -- 獲取目前 token_bucket 中剩餘 token 數 -- 如果是第一次進入，則設定 token_bucket 數量為 令牌桶最大值 local last_tokens = tonumber(redis.call("get", KEYS[1])) if last_tokens == nil then last_tokens = capacity end -- 上一次更新 token_bucket 的時間 local last_refreshed = tonumber(redis.call("get", KEYS[2])) if last_refreshed == nil then last_refreshed = 0 end local delta = math.max(0, now-last_refreshed) -- 通過當前時間與上一次更新時間的跨度，以及生產token的速率，計算出新的token數 -- 如果超過 max_burst，多餘生產的token會被丟棄 local filled_tokens = math.min(capacity, last_tokens+(delta*rate)) local allowed = filled_tokens >= requested local new_tokens = filled_tokens if allowed then new_tokens = filled_tokens - requested end -- 更新新的token數，以及更新時間 redis.call("setex", KEYS[1], ttl, new_tokens) redis.call("setex", KEYS[2], ttl, now) return allowed ``` 上述可以看出 `lua script` ：只涉及對 token 操作，保證 token 生產合理和讀取合理。 ## 函式分析  從上述流程中看出： 1. 有多重保障機制，保證限流一定會完成。 2. 如果`redis limiter`失效，至少在程序內`rate limiter`兜底。 3. 重試 `redis limiter` 機制保證儘可能地正常執行。 ## 總結 `go-zero` 中的 `tokenlimit` 限流方案適用於瞬時流量衝擊，現實請求場景並不以恆定的速率。令牌桶相當預請求，當真實的請求到達不至於瞬間被打垮。當流量衝擊到一定程度，則才會按照預定速率進行消費。 但是生產`token`上，不能按照當時的流量情況作出動態調整，不夠靈活，還可以進行進一步優化。此外可以參考[Token bucket WIKI](https://en.wikipedia.org/wiki/Token_bucket)中提到分層令牌桶，根據不同的流量頻寬，分至不同排隊中。 ## 參考 - [go-zero tokenlimit](https://github.com/tal-tech/go-zero/blob/master/core/limit/tokenlimit.go) - [Go-Redis 提供的分散式限流庫](https://github.com/go-redis/redis_rate) 如果覺得文章不錯，歡迎 github 點個star