golang版本使用令牌桶算法来实现限流的策略

代码

package main

import (
	"fmt"
	"sync"
	"time"
)

// 令牌桶结构体
type TokenBucket struct {
	tokens chan struct{}
	rate   time.Duration
	// 桶容量
	limit int
}

// 创建令牌桶
func NewTokenBucket(rate time.Duration, limit int) *TokenBucket {
	return &TokenBucket{
		tokens: make(chan struct{}, limit),
		rate:   rate,
		limit:  limit,
	}
}

// 放入令牌
func (t *TokenBucket) AddToken() {
	select {
	case t.tokens <- struct{}{}:
		{
			fmt.Println("放入成功")
		}
	default:
		// 桶已满，丢弃令牌
		fmt.Println("桶已满，丢弃令牌")
	}
}

// 取出令牌
func (t *TokenBucket) GetToken() bool {
	select {
	case <-t.tokens:
		return true
	case <-time.After(t.rate):
		// fmt.Println("获取失败")
		return false
	}
}

func main() {
	// 创建令牌桶，每秒放入 10 个令牌，桶容量为 10
	tb := NewTokenBucket(time.Second, 5)
	var success = make(chan struct{}, 3000)
	var fail = make(chan struct{}, 3000)
	var total