// WaitN 会发生阻塞直到 lim 允许的 n 个事件执行。 // 它返回一个 error 如果 n 超过了 Limiter的桶大小, // Context会被取消, 或等待的时间超过了 Context 的 Deadline。 func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)
// reserveN is a helper method for AllowN, ReserveN, and WaitN. // maxFutureReserve specifies the maximum reservation wait duration allowed. // reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN. 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, last, tokens := lim.advance(now)
// Calculate the remaining number of tokens resulting from the request. tokens -= float64(n)
// Calculate the wait duration var waitDuration time.Duration if tokens < 0 { waitDuration = lim.limit.durationFromTokens(-tokens) }
// Decide result ok := n <= lim.burst && waitDuration <= maxFutureReserve
// Prepare reservation r := Reservation{ ok: ok, lim: lim, limit: lim.limit, } if ok { r.tokens = n r.timeToAct = now.Add(waitDuration) }
// Update state if ok { lim.last = now lim.tokens = tokens lim.lastEvent = r.timeToAct } else { lim.last = last }
lim.mu.Unlock() return r }
// advance calculates and returns an updated state for lim resulting from the passage of time. // lim is not changed. func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) { last := lim.last if now.Before(last) { last = now }
// Avoid making delta overflow below when last is very old. maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens) elapsed := now.Sub(last) if elapsed > maxElapsed { elapsed = maxElapsed }
// Calculate the new number of tokens, due to time that passed. delta := lim.limit.tokensFromDuration(elapsed) tokens := lim.tokens + delta if burst := float64(lim.burst); tokens > burst { tokens = burst }
return now, last, tokens }
// durationFromTokens is a unit conversion function from the number of tokens to the duration // of time it takes to accumulate them at a rate of limit tokens per second. func (limit Limit) durationFromTokens(tokens float64) time.Duration { seconds := tokens / float64(limit) return time.Nanosecond * time.Duration(1e9*seconds) }
// tokensFromDuration is a unit conversion function from a time duration to the number of tokens // which could be accumulated during that duration at a rate of limit tokens per second. func (limit Limit) tokensFromDuration(d time.Duration) float64 { return d.Seconds() * float64(limit) }