十年网站开发经验 + 多家企业客户 + 靠谱的建站团队
量身定制 + 运营维护+专业推广+无忧售后,网站问题一站解决
起初是因为要去拉取一些第三方的数据,而第三方的API接口都有限流措施。比如6000/分钟,500/分钟。想着拉取数据就用多个协程的方式。但是容易超频,所以想着写一个限流的东东。网上有讲令牌桶类似下面这样:(网上的原理图)
成都网络公司-成都网站建设公司成都创新互联十年经验成就非凡,专业从事成都网站制作、做网站、外贸营销网站建设,成都网页设计,成都网页制作,软文推广,广告投放平台等。十年来已成功提供全面的成都网站建设方案,打造行业特色的成都网站建设案例,建站热线:18980820575,我们期待您的来电!
给人的感觉挺简单,于是来实践一下。首先来写桶的结构
这里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。