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

代码

package mainimport ("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 truecase <-time.After(t.rate):// fmt.Println("获取失败")return false}
}func main() {// 创建令牌桶，每秒放入 10 个令牌，桶容量为 10tb := NewTokenBucket(time.Second, 5)var success = make(chan struct{}, 3000)var fail = make(chan struct{}, 3000)var total = make(chan struct{}, 3000)var x intvar wg sync.WaitGroupwg.Add(1)go func() {defer func() {wg.Done()}()for i := 0; i < 11; i++ {// fmt.Println("尝试放入")tb.AddToken()time.Sleep(time.Millisecond * 500)tb.AddToken()time.Sleep(time.Millisecond * 500)}}()wg.Add(1)go func() {defer func() {wg.Done()}()for k := 0; k < 10; k++ {for j := 0; j < 100; j++ {wg.Add(1)go func(k, j int) {defer func() {wg.Done()total <- struct{}{}}()token := tb.GetToken()x++if token {success <- struct{}{}fmt.Printf("\033[32m请求通过![%d,%d]\033[0m\n", k, j)} else {fail <- struct{}{}fmt.Printf("\033[31m请求被拒绝![%d,%d]\033[0m\n", k, j)}}(k, j)}time.Sleep(time.Second)}}()wg.Wait()fmt.Printf("成功: %d, 失败: %d，总数：%d\n", len(success), len(fail), len(total))fmt.Println("end", x)}