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

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
}

package tokenBucket
import (
   "log"
   "sync"
   "time"
)
type TokensBucket struct {
   limiter float64    //速率
   burst   int        //桶大小
   mu      sync.Mutex //锁
   tokens  float64    //桶里面的令牌数量
   last    time.Time  //最后一次消耗令牌的时间
}
// NewTokensBucket 创建令牌桶
func NewTokensBucket(limiter float64, burst int) *TokensBucket {
   return &TokensBucket{limiter: limiter, burst: burst}
}
// Allow 使用，每次消耗一个令牌
func (t *TokensBucket) Allow() bool {
   return t.AllowN(time.Now(), 1)
}
// AllowN 当前时间，一次消耗的令牌
func (t *TokensBucket) AllowN(now time.Time, i int) bool {
   t.mu.Lock()
   defer t.mu.Unlock()
   //当前时间-最后一次添加令牌的时间 * 桶速率 = 应该补充的令牌
   delta := now.Sub(t.last).Seconds() * t.limiter
   t.tokens += delta
   //桶内令牌 > 桶总大小  =  只补充最大令牌数
   if t.tokens > float64(t.burst) {
      t.tokens = float64(t.burst)
   }
   //桶内令牌 < 需要的令牌 = 返回false
   if t.tokens < float64(i) {
      return false
   }
   //否则返回true，并用桶的剩余令牌 - 消耗令牌
   t.tokens -= float64(i)
   //桶最后一次补充时间重置为当前时间
   t.last = now
   //返回true
   return true
}

func main() {
   bucket := NewTokensBucket(3, 5)
   for true {
      n := 4
      for i := 0; i < n; i++ {
         go func(i int) {
            if bucket.Allow() {
               log.Printf("allow [%d]", i)
            } else {
               log.Printf("forbid [%d]", i)
            }
         }(i)
      }
      time.Sleep(time.Second)
      log.Println("========================================")
   }
}

func main() {
    app := gin.Default()
    bucket := tokenBucket.NewTokensBucket(1, 2)
    app.Use(func(context *gin.Context) {
   //拿到令牌就给放行
        if bucket.Allow() {
            context.Next()
    //拿不到就不给过
        } else {
            context.JSON(500, gin.H{
                "msg": "false",
            })
            context.Abort()
        }
    })
    app.GET("/", func(context *gin.Context) {
        context.JSON(200, gin.H{
            "msg": "success",
        })
    })
    app.Run(":80")
}