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

在上一篇文章 Go微服务: 令牌桶 当中简单的介绍了令牌桶实现的原理，然后利用 /x/time/rate 这个库 10 行代码写了一个基于 ip 的 gin 限流中间件，那这个功能是怎么实现的呢？接下来我们就从源码层面来了解一下这个库的实现。这个实现很有意思，并没有真正的使用一个定时器不断的生成令牌，而是靠计算的方式来完成

2.rate/limt

在golang.org/x/time/rate库中
使用限速器的时候我们需要调用 NewLimiter 方法，然后 Limiter 提供了三组限速的方法，这三组方法其实都是通过调用 reserveN 实现的 reserveN 返回一个 *Reservation 指针，先来看一下这两个结构体。

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

预定，表示预约某个时间的 token

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

这个库并没有使用定时器来发放 token 而是用了 lazyload 的方式，等需要消费 token 的时候才通过时间去计算然后更新 token 的数量，下面我们先通过一个例子来看一下这个流程是怎么跑的

如上图所示，假设我们有一个限速器，它的 token 生成速度为 1，也就是一秒一个，桶的大小为 10，每个格子表示一秒的时间间隔

last表示上一次更新 token时还有 2 个token

现在我们有一个请求竟来， 总共需要7个 token才能完成请求

now表示我现在进来的时间，距离last 已经过去了2s, 那么现在就有4个token(每秒生成一个token)

所以,如果需要 7 个 token 那么也就还需要等待 3s 中才真的有 7 个，所以这就是 timeToAct 所在的时间节点

预约成功之后更新 last = now 、token = -3 因为 token 已经被预约出去了所以现在剩下的就是负数了

2.3 消费 token

总共有三种消费 token 的方法 AllowN, ReserveN, WaitN最终都是调用的reserveN 这个方法

// 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 }

advance 方法用于计算 token 的数量

// 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 }