QPS限流之令牌桶

type LingLimit struct {
	// 锁
	sync.Mutex
	// 模拟时钟
	clock Clock
	// 开始取令牌的第一个时间
	startTime time.Time
	// 容量
	capacity int64
	// 给桶中加令牌的数量（每一时刻）
	quantum int64
	// 每个时刻之间的间隔大小
	fillInterval time.Duration
	// 桶中可用的令牌数量
	availableToken int64
	// 上一时刻
	latestTick int64
}

// 默认最大等待时间
const InfinityDuration time.Duration = 0x7fffffffffffffff

func NewLingLimit(fillInterval time.Duration, capacity int64, quantum int64, clock Clock) *LingLimit {
	if clock == nil {
		clock = realClock{}
	}
	if fillInterval <= 0 {
		panic("token bucket fill interval is not > 0")
	}
	if capacity <= 0 {
		panic("token bucket capacity is not > 0")
	}
	if quantum <= 0 {
		panic("token bucket quantum is not > 0")
	}

	return &LingLimit{
		clock:        clock,
		startTime:    clock.Now(),
		capacity:     capacity,
		quantum:      quantum,
		fillInterval: fillInterval,
		// 一开始可用令牌数为初始容量
		availableToken: capacity,
		latestTick:     0,
	}
}

func (o *LingLimit) take(now time.Time, count int64, maxWait time.Duration) (time.Duration, bool) {
	o.Lock()
	defer o.Unlock()
	// 有多少个时刻
	tick := o.currentTick(now)
	// 调整可以使用的令牌数量
	o.adjustavailableTokens(tick)

	avail := o.availableToken - count
	if avail >= 0 {
		// 当前剩余令牌
		o.availableToken = avail
		// 返回可以获得令牌
		return 0, true
	}

	// 需要更长的时刻来累计令牌数量
	endTick := tick + (-avail + o.quantum - 1) / o.quantum
	// 最终能拿到相应数量的令牌的时间
	endTime := o.startTime.Add(time.Duration(endTick) * o.fillInterval)
	// 需要等待的时间
	waitTime := endTime.Sub(now)
	if waitTime > maxWait{
		return 0, false
	}

	// 可用的令牌数是一个负值，表示超额使用
	o.availableToken = avail

	return waitTime, true

}

func (o *LingLimit) currentTick(now time.Time) int64 {
	return int64(now.Sub(o.startTime) / o.fillInterval)
}

// 调整可用的令牌数
func (o *LingLimit) adjustavailableTokens(tick int64) {
	lastTick := o.latestTick
	o.latestTick = tick

	// 如果当前时间内的可用令牌数大于最大容量
	if o.availableToken >= o.capacity {
		return
	}

	// 这一时刻，减去上一时刻（中间间隔多少时刻）* 每一个时刻放入多少令牌 = 当前可用令牌
	o.availableToken += (tick - lastTick) * o.quantum
	if o.availableToken >= o.capacity {
		o.availableToken = o.capacity
		return
	}

	return
}

模拟时钟工具类

// Clock represents the passage of time in a way that
// can be faked out for tests.
type Clock interface {
	// Now returns the current time.
	Now() time.Time
	// Sleep sleeps for at least the given duration.
	Sleep(d time.Duration)
}

// realClock implements Clock in terms of standard time functions.
type realClock struct{}

// Now implements Clock.Now by calling time.Now.
func (realClock) Now() time.Time {
	return time.Now()
}

// Now implements Clock.Sleep by calling time.Sleep.
func (realClock) Sleep(d time.Duration) {
	time.Sleep(d)
}

时钟可以使用工具类

"github.com/andres-erbsen/clock"