粗略的写个令牌桶算法 – 兰陵美酒郁金香的个人博客

如下

package main

import (
	"fmt"
	"time"
)

// Bucket 令牌桶的抽象实例
type Bucket struct {
	// bucket 为桶队列, stopSignal 作为停止信号
	bucket, stopSignal chan struct{}
	// 桶的令牌数
	num                int
	// 单位时间
	duration           time.Duration
}

// New 构造初始化
func New(num int, t time.Duration) *Bucket {
	b := &Bucket{
		bucket:     make(chan struct{}, num),
		stopSignal: make(chan struct{}),
		num:        num,
		duration:   t,
	}
	go b.run()
	return b
}

// run 启动函数
func (b *Bucket) run() {
	// 先初始化给塞满桶
	for i := 0; i < b.num; i++ {
		b.bucket <- struct{}{}
	}
	// 再来个定时器
	ticker := time.NewTicker(b.duration)
	defer ticker.Stop()
	defer close(b.bucket)
FOR:
	for {
		select {
		case <-ticker.C:
			// 如果已经满了,为防止超限,还要等待一会儿
			if len(b.bucket) == b.num {
				time.Sleep(b.duration)
				continue
			}
			// 我们需要补充桶的令牌数
			count:=b.num-len(b.bucket)
			for i := 0; i < count; i++ {
				b.bucket <- struct{}{}
			}
		case <-b.stopSignal:
			break FOR
		}
	}
}

// Stop 发送停止信号
func (b *Bucket) Stop() {
	b.stopSignal <- struct{}{}
}

// Get 获取一个令牌,如果没拿到,返回 false
func (b *Bucket) Get() bool {
	select {
	case <-b.bucket:
		return true
	default:
		return false
	}
}

func main() {
	// 初始化一秒钟 5 个令牌
	b := New(5, time.Second)
	// 模拟 10 个人抢
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
				time.Sleep(time.Second / 2)
				if !b.Get() {
					fmt.Printf("协程 %d没抢到\n", i)
					continue
				}
				fmt.Printf("协程 %d 抢到了\n", i)
			}
		}(i)
	}

	select {}

}

package main

import (

"fmt"

"time"

)

// Bucket 令牌桶的抽象实例

type Bucket struct {

// bucket 为桶队列, stopSignal 作为停止信号

bucket, stopSignal chan struct{}

// 桶的令牌数

num int

// 单位时间

duration time.Duration

}

// New 构造初始化

func New(num int, t time.Duration) *Bucket {

b := &Bucket{

bucket: make(chan struct{}, num),

stopSignal: make(chan struct{}),

num: num,

duration: t,

}

go b.run()

return b

}

// run 启动函数

func (b *Bucket) run() {

// 先初始化给塞满桶

for i := 0; i < b.num; i++ {

b.bucket <- struct{}{}

}

// 再来个定时器

ticker := time.NewTicker(b.duration)

defer ticker.Stop()

defer close(b.bucket)

FOR:

for {

select {

case <-ticker.C:

// 如果已经满了,为防止超限,还要等待一会儿

if len(b.bucket) == b.num {

time.Sleep(b.duration)

continue

}

// 我们需要补充桶的令牌数

count:=b.num-len(b.bucket)

for i := 0; i < count; i++ {

b.bucket <- struct{}{}

}

case <-b.stopSignal:

break FOR

}

// Stop 发送停止信号

func (b *Bucket) Stop() {

b.stopSignal <- struct{}{}

}

// Get 获取一个令牌,如果没拿到,返回 false

func (b *Bucket) Get() bool {

select {

case <-b.bucket:

return true

default:

return false

}

func main() {

// 初始化一秒钟 5 个令牌

b := New(5, time.Second)

// 模拟 10 个人抢

for i := 0; i < 10; i++ {

go func(i int) {

for {

time.Sleep(time.Second / 2)

if !b.Get() {

fmt.Printf("协程 %d没抢到\n", i)

continue

}

fmt.Printf("协程 %d 抢到了\n", i)

}

}(i)

}

select {}

}

暂无

兰陵美酒郁金香

大道至简 Simplicity is the ultimate form of sophistication.

文章评论（0）

你必须登录才能发表评论