起初是因為要去拉取一些第三方的資料，而第三方的API介面都有限流措施，比如6000/分鐘，500/分鐘，想著拉取資料就用多個協程的方式，但是容易超頻，所以想著寫一個限流的東東，網上有講令牌桶類似下面這樣：(網上的原理圖)
令牌桶原理
- 有一個桶，桶有容量(cap：桶的容量)，
- 然后以恒定的速度往桶里加入令牌(token:表示令牌)，
- 如果桶已經達到容量，新加入的令牌將被廢棄，
- 每次消耗就是從桶里拿走一個令牌，
給人的感覺挺簡單，于是來實踐一下，首先來寫桶的結構
桶結構體
這里token是存放令牌的地方，這里用了一個空的struct{}，眾所周知空struct{}耗記憶體極少，mu是一個互斥鎖，因為涉及到多個協程操作到桶內的令牌，所以這里加了一個鎖，
package limit
import (
"sync"
"time"
)
// 令牌桶
type bucket struct {
rate int // 每分鐘頻率(每分鐘加入多少個令牌)
token chan struct{} // 存放令牌的地方
cap int // 容量
mu *sync.Mutex // 桶內的鎖
pause bool // 暫停
stop bool // 停止
}
實體化一個桶
這里判斷了一下桶的容量必須大于0，
// 獲取新的bucket
// rate: 每分鐘多少次
// cap: 桶的容量,必須大于等于1
func NewBucket(rate, cap int) *bucket {
if cap < 1 {
panic("limit bucket cap error")
}
return &bucket{
token: make(chan struct{}, cap),
rate: rate,
mu: new(sync.Mutex),
cap: cap,
}
}
開始計時
這里用一個新的goroutine來計時
// 開始
func (b *bucket) Start() {
go b.addToken()
}
// 加入令牌
func (b *bucket) addToken() {
for {
b.mu.Lock()
if b.stop {
close(b.token)
b.mu.Unlock()
return
}
if b.pause {
b.mu.Unlock()
time.Sleep(time.Second)
continue
}
b.token <- struct{}{}
d := time.Minute / time.Duration(b.rate)
b.mu.Unlock()
time.Sleep(d)
}
}
消費一個令牌
這里獲取一個令牌就是從chan里拿一個資料
// 消費,這里會自動阻塞
func (b *bucket) GetToken() {
<-b.token
}
暫停，停止，重置
利用桶的屬性：pause，stop 還可以加一些小的功能，
// 暫停
func (b *bucket) Pause() {
b.mu.Lock()
defer b.mu.Unlock()
b.pause = true
}
// 停止
func (b *bucket) Stop() {
b.mu.Lock()
defer b.mu.Unlock()
b.stop = true
}
// 重置
func (b *bucket) Reset() {
b.mu.Lock()
defer b.mu.Unlock()
b.token = make(chan struct{}, b.cap)
}
基本已經實作我們想要的功能，測驗一下，這時同事甩給我一個包：golang.org/x/time/rate，
轉載請註明出處，本文鏈接：https://www.uj5u.com/houduan/436301.html
標籤：Go