使用 Golang 实现简易的令牌桶算法

安正勋

已于 2024-08-24 16:46:49 修改

阅读量2.3k

点赞数

分类专栏：朝花夕拾文章标签：令牌桶算法漏桶算法简易实现 Golang

于 2022-04-24 20:31:49 首次发布

本文链接：https://blog.csdn.net/qq_35246620/article/details/124390385

版权

朝花夕拾专栏收录该内容

38 篇文章 5 订阅

订阅专栏

文章目录

简介
实现
轮子

简介

在网络中传输数据的时候时，为了防止网络拥塞，需限制流出网络的流量，使流量以比较均匀的速度向外发送。令牌桶算法就实现了这个功能，可控制发送到网络上数据的数目，并允许突发数据的发送。

令牌桶算法是网络流量整形和速率限制中最常使用的一种算法。大小固定的令牌桶可自行以恒定的速率源源不断地产生令牌。如果令牌不被消耗，或者被消耗的速度小于产生的速度，令牌就会不断地增多，直到把桶填满。后面再产生的令牌就会从桶中溢出。最后桶中可以保存的最大令牌数永远不会超过桶的大小。
ipt-1
传送到令牌桶的数据包需要消耗令牌。不同大小的数据包，消耗的令牌数量不一样。

令牌桶这种控制机制基于令牌桶中是否存在令牌来指示什么时候可以发送流量。令牌桶中的每一个令牌都代表一个字节。如果令牌桶中存在令牌，则允许发送流量；而如果令牌桶中不存在令牌，则不允许发送流量。因此，如果突发门限被合理地配置并且令牌桶中有足够的令牌，那么流量就可以以峰值速率发送。

lpt-2
与“令牌桶算法”类似的算法还有“漏桶算法”，这两种算法的主要区别在于“漏桶算法”能够强行限制数据的传输速率，而“令牌桶算法”在能够限制数据的平均传输速率外，还允许某种程度的突发传输。在“令牌桶算法”中，只要令牌桶中存在令牌，那么就允许突发地传输数据直到达到用户配置的门限，因此它适合于具有突发特性的流量。

在本文中，我们使用 Golong 语言实现一个简单的“令牌桶算法”。

实现

首先，我们假设令牌桶的放入令牌的速率是恒定的，不考虑流量速率突变的情况。

package awesomeProject

import (
	"sync"
	"time"
)

// 定义令牌桶结构
type tokenBucket struct {
	limitRate int           // 限制频率，即每分钟加入多少个令牌
	tokenChan chan struct{} // 令牌通道，可以理解为桶
	cap       int           // 令牌桶的容量
	muLock    *sync.Mutex   // 令牌桶锁，保证线程安全
	stop      bool          // 停止标记，结束令牌桶
}

// NewTokenBucket 创建令牌桶
func NewTokenBucket(limitRate, cap int) *tokenBucket {
	if cap < 1 {
		panic("token bucket cap must be large 1")
	}
	return &tokenBucket{
		tokenChan: make(chan struct{}, cap),
		limitRate: limitRate,
		muLock:    new(sync.Mutex),
		cap:       cap,
	}
}

// Start 开启令牌桶
func (b *tokenBucket) Start() {
	go b.produce()
}

// 生产令牌
func (b *tokenBucket) produce() {
	for {
		b.muLock.Lock()
		if b.stop {
			close(b.tokenChan)
			b.muLock.Unlock()
			return
		}
		b.tokenChan <- struct{}{}
		d := time.Minute / time.Duration(b.limitRate)
		b.muLock.Unlock()
		time.Sleep(d)
	}
}

// Consume 消费令牌
func (b *tokenBucket) Consume() {
	<-b.tokenChan
}

// Stop 停止令牌桶
func (b *tokenBucket) Stop() {
	b.muLock.Lock()
	defer b.muLock.Unlock()
	b.stop = true
}