微服務-限流：一.golang實現令牌桶算法

起初是因為要去拉取一些第三方的數據，而第三方的API接口都有限流措施。比如6000/分鍾，500/分鍾。想着拉取數據就用多個協程的方式。但是容易超頻，所以想着寫一個限流的東東。網上有講令牌桶類似下面這樣：(網上的原理圖)

令牌桶原理

1. 有一個桶，桶有容量(cap：桶的容量)。
2. 然后以恆定的速度往桶里加入令牌(token:表示令牌)。
3. 如果桶已經達到容量，新加入的令牌將被廢棄。
4. 每次消耗就是從桶里拿走一個令牌。

給人的感覺挺簡單，於是來實踐一下。首先來寫桶的結構

桶結構體

這里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。