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 = make(chan struct{}, 3000)
	var x int
	var wg sync.WaitGroup
	wg.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)

}

注意点

我在进行测试的时候，一开始success和fail这两个变量使用的是切片，我在多次运行后发现，succsss+fail并不等于total，但是打印的执行流程是没有问题的，于是想到切片是存在数据并发安全问题的，改为channel才得到了正确的结果