百木园-与人分享,
就是让自己快乐。

微服务-限流:一.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。


来源:https://www.cnblogs.com/ourongxin/p/15957811.html
本站部分图文来源于网络,如有侵权请联系删除。

未经允许不得转载:百木园 » 微服务-限流:一.golang实现令牌桶算法

相关推荐

  • 暂无文章