令牌桶算法的基本过程如下：
- 假如用户配置的平均发送速率为r，则每隔1/r秒一个令牌被加入到桶中；
- 假设桶最多可以存发b个令牌。如果令牌到达时令牌桶已经满了，那么这个令牌会被丢弃；
- 当一个n个字节的数据包到达时，就从令牌桶中删除n个令牌，并且数据包被发送到网络；
package main
import (
"fmt"
"sync"
"time"
)
type Bucket struct {
cap int //容量
ch chan bool
timer *time.Ticker //定时填充token
mu sync.Mutex
}
func NewBucket(cap int, interval time.Duration) *Bucket {
bucket := &Bucket{
cap: cap,
ch: make(chan bool, cap),
timer: time.NewTicker(interval),
}
go bucket.startTicker()
return bucket
}
func (bucket *Bucket) startTicker() {
for i := 0; i < bucket.cap; i++ {
bucket.ch <- true
}
for {
select {
case <-bucket.timer.C:
for i := len(bucket.ch); i < bucket.cap; i++ {
bucket.Add()
}
}
}
}
func (bucket *Bucket) Add() {
bucket.mu.Lock()
defer bucket.mu.Unlock()
if len(bucket.ch) < bucket.cap {
bucket.ch <- true
}
}
func (bucket *Bucket) Get() bool {
select {
case <-bucket.ch:
return true
default:
return false
}
}
func main() {
bucket := NewBucket(5, time.Second)
for i := 0; i < 1000; i++ {
time.Sleep(time.Millisecond * 100)
go DoFunc(bucket, i)
}
for {
}
}
func DoFunc(bucket *Bucket, index int) {
if bucket.Get() {
fmt.Printf("######### success : %d \n", index)
bucket.Add() // 每次请求成功后执行Add是否合理？？？每一时刻最多可有cap个token可用
} else {
fmt.Printf("######### failed : %d \n", index)
}
}
