Go微服务--令牌桶实现原理

type Limiter struct {    // 互斥锁	mu     sync.Mutex     // 每秒产生 token 的速度, 其实是 float64 的一个别名	limit  Limit     // 桶的大小	burst  int     // 当前时间节点拥有的 tokens 数量	tokens float64 	// 上次更新 token 的时间	last time.Time 	// 上次限速的时间，这个时间可能是过去的某个时间也可能是将来的某个时间	lastEvent time.Time}

type Reservation struct {    // 是否能预约上	ok        bool    // limter	lim       *Limiter    // 预约的 token 数量	tokens    int    // token 实际使用的时间	timeToAct time.Time	// 保存一下速率，因为 lim 的速率是可以被动态调整的，所以不能直接用	limit Limit}

// now: 需要消费 token 的时间点// n: 需要多少个 token// maxFutureReserve: 能够等待的最长时间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,		}	}     // advance 方法会去计算当前有多少个 token    // 后面会讲到，now 其实就是传入的时间，但是 last 可能会变	now, last, tokens := lim.advance(now) 	// 发放 token 之后还剩多少	tokens -= float64(n) 	// 根据 token 数量计算需要等待的时间	var waitDuration time.Duration	if tokens < 0 {		waitDuration = lim.limit.durationFromTokens(-tokens)	} 	// 计算是否可以发放，如果需要的量比桶的容量还大肯定是不行的    // 然后就是看需要能否容忍需要等待的时间	ok := n <= lim.burst && waitDuration <= maxFutureReserve 	// Prepare reservation	r := Reservation{		ok:    ok,		lim:   lim,		limit: lim.limit,	}    // 如果可以的话，就把 token 分配给预约者	if ok {		r.tokens = n		r.timeToAct = now.Add(waitDuration)	} 	// 更新各个字段的状态	if ok {		lim.last = now		lim.tokens = tokens		lim.lastEvent = r.timeToAct	} else {        // 为什么不 ok 也要更新 last 呢？因为 last 可能会改变		lim.last = last	} 	lim.mu.Unlock()	return r}

// now 是传入的当前的时间点，返回的 newNow 其实就是传入的参数，没有任何改变// newLast 是更新 token 的时间// newTokens 是 token 的数量func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {	// 如果当前时间比上次更新 token 的时间还要早，那么就重置一下 last    last := lim.last	if now.Before(last) {		last = now	} 	// 这里为了防止溢出，先计算了将桶填满需要花费的最大时间	maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)    // 计算时间差，如果大于最大时间的话，就取最大值	elapsed := now.Sub(last)	if elapsed > maxElapsed {		elapsed = maxElapsed	} 	// 计算这段时间生成的 token 数量，如果大于桶的容量，就取桶的容量	delta := lim.limit.tokensFromDuration(elapsed)	tokens := lim.tokens + delta	if burst := float64(lim.burst); tokens > burst {		tokens = burst	} 	return now, last, tokens}

func (limit Limit) durationFromTokens(tokens float64) time.Duration {	seconds := tokens / float64(limit)	return time.Nanosecond * time.Duration(1e9*seconds)}

func (limit Limit) tokensFromDuration(d time.Duration) float64 {	// 这里通过拆分整数和小数部分可以减少时间上的误差	sec := float64(d/time.Second) * float64(limit)	nsec := float64(d%time.Second) * float64(limit)	return sec + nsec/1e9}

// ctx 用于控制超时, n 是需要消费的 token 数量，如果 context 的 Deadline 早于要等待的时间就会直接返回失败func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {	lim.mu.Lock()	burst := lim.burst	limit := lim.limit	lim.mu.Unlock()     // 先看一下是不是已经超出消费极限了	if n > burst && limit != Inf {		return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst)	}     // 如果 ctx 已经结束了也不用等了	select {	case <-ctx.Done():		return ctx.Err()	default:	} 	// 计算一下可以等待的时间	now := time.Now()	waitLimit := InfDuration	if deadline, ok := ctx.Deadline(); ok {		waitLimit = deadline.Sub(now)	} 	// 调用 reserveN 得到预约数据	r := lim.reserveN(now, n, waitLimit)     // 如果不 ok 说明预约不到	if !r.ok {		return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)	} 	// 如果可以预约到，计算一下需要等多久	delay := r.DelayFrom(now)	if delay == 0 {		return nil	}     // 启动一个 timer 进行定时	t := time.NewTimer(delay)	defer t.Stop()	select {	case <-t.C:		// We can proceed.		return nil	case <-ctx.Done():		// 如果 context 主动取消了，那么之前预约的 token 数量需要归还		r.Cancel()		return ctx.Err()	}}

func (r *Reservation) CancelAt(now time.Time) {    // 不 ok 说明没有预约上，直接返回就行了	if !r.ok {		return	} 	r.lim.mu.Lock()	defer r.lim.mu.Unlock()     // 如果没有速率限制，或者没有消费 token 或 token 已经被消费了，都不用还了	if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {		return	} 	// 计算需要还的 token 数量    // 这里说是需要减去已经预支的 token 数量，但是我发现应该是个 bug，感觉这里减重复了	restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))	if restoreTokens <= 0 {		return	}     // 计算当前拥有的 tokens 数量	now, _, tokens := r.lim.advance(now) 	// 当前拥有的加上需要归还的就是现有的，但是不能大于桶的容量	tokens += restoreTokens	if burst := float64(r.lim.burst); tokens > burst {		tokens = burst	} 	// 更新 tokens 数量	r.lim.last = now	r.lim.tokens = tokens     // 如果相等说明后面没有新的 token 消费，所以将状态重置到上一次	if r.timeToAct == r.lim.lastEvent {		prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))		if !prevEvent.Before(now) {			r.lim.lastEvent = prevEvent		}	} 	return}

func wait() {	l := rate.NewLimiter(10, 10)	t := time.Now()	l.ReserveN(t, 10) 	var wg sync.WaitGroup 	ctx, cancel := context.WithTimeout(context.TODO(), time.Hour)	defer cancel()     // 注释掉下面这段就不会提前 cancel	wg.Add(1)	go func() {		defer wg.Done()		// 模拟出现问题, 200ms就取消了		time.Sleep(200 * time.Millisecond)		cancel()	}() 	wg.Add(2)	go func() {		defer wg.Done()		// 如果要等，这个要等 1s 才能执行，但是我们的 ctx 200ms 就会取消		l.WaitN(ctx, 10)		fmt.Printf("[1] cost: %s\n", time.Since(t))	}() 	time.Sleep(100 * time.Millisecond) 	go func() {		defer wg.Done()		// 正常情况下，这个要等 1.2 s 才能执行，但是我们前面都取消了		// 这个是不是应该就只需要等 200ms 就执行了		ctx, cancel := context.WithTimeout(context.Background(), time.Hour)		defer cancel()		l.WaitN(ctx, 2)		fmt.Printf("[2] cost: %s\n", time.Since(t))	}() 	wg.Wait()}

[1] cost: 1.0002113s[2] cost: 1.2007347s

[1] cost: 200.8268ms[2] cost: 1.201066s

restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))

func main() {	t0 := time.Now()	t1 := time.Now().Add(100 * time.Millisecond)	t2 := time.Now().Add(200 * time.Millisecond)	t3 := time.Now().Add(300 * time.Millisecond) 	l := rate.NewLimiter(10, 20)	l.ReserveN(t0, 15) // 桶里还剩 5 个 token	fmt.Printf("%+v\n", l) 	r := l.ReserveN(t1, 10) // 桶还有 -4 个，	fmt.Printf("%+v\n", l) 	// 注释掉下面两行，最后结果还剩 8 个 token	l.ReserveN(t2, 2) // 桶里还有 -5 个	fmt.Printf("%+v\n", l) 	r.CancelAt(t3)	fmt.Printf("%+v\n", l)	// 归还之前借的，运行结果 桶里还有 4 个	// 但是这里不应该剩下 6 个么，本来有 5 个，300ms 生成了 3 个，后面又预支出去 2 个	// 而且我发现如果我注释掉预支两个的代码，结果和我预期的一致，剩余 8 个token}