Golang限流器time/rate设计与实现详解

type Limiter struct {
	mu        sync.Mutex // 互斥锁（排他锁）
	limit     Limit      // 放入桶的频率  float64 类型
	burst     int        // 桶的大小
	tokens    float64    // 令牌 token 当前剩余的数量
	last      time.Time  // 最近取走 token 的时间
	lastEvent time.Time  // 最近限流事件的时间
}

type Reservation struct {
	ok        bool      // 是否满足条件分配了 token
	lim       *Limiter  // 发送令牌的限流器
	tokens    int       // 发送 token 令牌的数量
	timeToAct time.Time // 满足令牌发放的时间
	limit     Limit     // 令牌发放速度
}

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool 

func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
	lim.mu.Lock()

	// 首先判断是否放入频率是否为无穷大
	// 如果为无穷大，说明暂时不限流
	if lim.limit == Inf {
		lim.mu.Unlock()
		return Reservation{
			ok:        true,
			lim:       lim,
			tokens:    n,
			timeToAct: now,
		}
	}

	// 拿到截至 now 时间时
	// 可以获取的令牌 tokens 数量及上一次拿走令牌的时间 last
	now, last, tokens := lim.advance(now)

	// 更新 tokens 数量
	tokens -= float64(n)

	// 如果 tokens 为负数，代表当前没有 token 放入桶中
	// 说明需要等待，计算等待的时间
	var waitDuration time.Duration
	if tokens < 0 {
		waitDuration = lim.limit.durationFromTokens(-tokens)
	}

	// 计算是否满足分配条件
	// 1、需要分配的大小不超过桶的大小
	// 2、等待时间不超过设定的等待时长
	ok := n <= lim.burst && waitDuration <= maxFutureReserve

	// 预处理 reservation
	r := Reservation{
		ok:    ok,
		lim:   lim,
		limit: lim.limit,
	}
	// 若当前满足分配条件
	// 1、设置分配大小
	// 2、满足令牌发放的时间 = 当前时间 + 等待时长
	if ok {
		r.tokens = n
		r.timeToAct = now.Add(waitDuration)
	}

	// 更新 limiter 的值，并返回
	if ok {
		lim.last = now
		lim.tokens = tokens
		lim.lastEvent = r.timeToAct
	} else {
		lim.last = last
	}

	lim.mu.Unlock()
	return r
}

func NewLimiter(r Limit, b int) *Limiter {
	return &Limiter{
		limit: r, // 放入桶的频率
		burst: b, // 桶的大小
	}
}

func main() {
	r := rate.Every(1 * time.Millisecond)
	limit := rate.NewLimiter(r, 10)
	http.HandleFunc("/", func(writer http.ResponseWriter, request *http.Request) {
		if limit.Allow() {
			fmt.Printf("请求成功，当前时间：%s\n", time.Now().Format("2006-01-02 15:04:05"))
		} else {
			fmt.Printf("请求成功，但是被限流了。。。\n")
		}
	})

	_ = http.ListenAndServe(":8081", nil)
}

func GetApi() {
	api := "http://localhost:8081/"
	res, err := http.Get(api)
	if err != nil {
		panic(err)
	}
	defer res.Body.Close()

	if res.StatusCode == http.StatusOK {
		fmt.Printf("get api success\n")
	}
}

func Benchmark_Main(b *testing.B) {
	for i := 0; i < b.N; i++ {
		GetApi()
	}
}