rate.Limiter原理剖析

rate.limiter，是golang官方提供的限流器，采用了令牌算法。对限流算法陌生的读者朋友，可以参考我的另一篇文章服务器限流算法与实现

基础用法

官方实现的limiter提供了丰富的功能，包括：

消耗一个token
消耗N个token
等待获取N个token
预定N个token，并可取消预定

package main

import (
    "fmt"
    "time"

    "golang.org/x/time/rate"
)

func main() {
    limit := rate.Every(time.Millisecond * 10)

    limiter := rate.NewLimiter(limit, 5)

    for i := 0; i < 100; i++ {
        if limiter.Allow() {
            println("allow", i)
        } else {
            println("limit", i)
        }
        time.Sleep(time.Millisecond * 5)
    }
    fmt.Println("done")

}

核心源码解析

limiter.drawio.png

limiter属性与方法

limiter数据结构比较简单：

mu：一个锁
limit：每秒允许的token数
burst：允许缓存的最大token数
last：token上次更新的时间
lastEvent：最新一次处理的limit事件时间，可能会早于晚于当前时间

limiter支持如下方法使用：

Tokens()：获取当前可用的token数
Allow() & AllowN()：如果token足够，消耗指定数量token，否则返回false
Reserve() & ReserveN()：预定指定数量的token，同时消耗token额度，否则返回false
Wait() & WaitN()：等待获取指定数量的token，否则返回false

Reservation，当limiter对象调用Reserve()方法时，会获得一个Reservation对象，该对象数据结构包括：

ok：是否预定成功
lim：关联的limiter对象
tokens：预定的token数
timeToAct：预定token的获取时间

Reservation提供了获取时间差、取消预定的方法：

Delay() & DelayFrom()：预定时间距指定时间的时间差
Cancel() & CancelAt()：在指定时间取消预定，同时会返还有效的token到limiter

limiter主函数reserveN()分析

limiter开放的allow、reserve、wait方法，都会调用内部的reserveN()方法，该方法入参为期望时间、期望数量、最大等待时间。

// 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(t time.Time, n int, maxFutureReserve time.Duration) Reservation {}

首先获取mu.Lock()，避免并发修改
然后判断limit限制数是否为边界值，如果无限制，则返回true，如果limit为0，则判断最大容量burst值是否满足需要
计算t时间时，可用的token数

func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration) Reservation {
    ...
    t, tokens := lim.advance(t)    
    ...
}


func (lim *Limiter) advance(t time.Time) (newT time.Time, newTokens float64) {
    last := lim.last
    if t.Before(last) {
        last = t
    }

    // Calculate the new number of tokens, due to time that passed.
    elapsed := t.Sub(last)
    delta := lim.limit.tokensFromDuration(elapsed)
    tokens := lim.tokens + delta
    if burst := float64(lim.burst); tokens > burst {
        tokens = burst
    }
    return t, tokens
}

4. 计算可用的token数能否满足需要，如果不够，需要等待的时间，是否小于最大等待时间maxFutureReserve

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

5. 如果满足需求，则更新reservation，更新limit的状态。其中，last为当前获取的时间，lastEvent，为此次预定的执行时间

    if ok {
        r.tokens = n
        r.timeToAct = t.Add(waitDuration)

        // Update state
        lim.last = t
        lim.tokens = tokens
        lim.lastEvent = r.timeToAct
    }

reservation函数cancel()分析

当持有reservation的对象不在执行，可以调用cancel方法，将获取的token返回给limiter。具体返回的token数量，根据cancel的时间来计算。

首先判断reasevation是否有效，timeToAct是否在取消的时间点之后
根据timeToAct、limiter的lastEvent，计算可以归还的token数量

    // calculate tokens to restore
    // The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved
    // after r was obtained. These tokens should not be restored.
    restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
    if restoreTokens <= 0 {
        return
    }

3. 计算到指定时间，产生的token数，更新token数量

    // advance time to now
    t, tokens := r.lim.advance(t)
    // calculate new number of tokens
    tokens += restoreTokens
    if burst := float64(r.lim.burst); tokens > burst {
        tokens = burst
    }

4. 更新limiter的token数，last、lastEvent

    // update state
    r.lim.last = t
    r.lim.tokens = tokens
    if r.timeToAct == r.lim.lastEvent {
        prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
        if !prevEvent.Before(t) {
            r.lim.lastEvent = prevEvent
        }
    }

limiter函数wait()分析

相比较于allow，wait方法多了定时器功能，监听context的超时时间，如果获得reservation的timeToAct早于超时时间，则返回成功，否则调用reservation的cancel方法，取消预定，wait返回失败

// Determine wait limit
    waitLimit := InfDuration
    if deadline, ok := ctx.Deadline(); ok {
        waitLimit = deadline.Sub(t)
    }
    // Reserve
    r := lim.reserveN(t, n, waitLimit)
    if !r.ok {
        return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
    }
    // Wait if necessary
    delay := r.DelayFrom(t)
    if delay == 0 {
        return nil
    }
    ch, stop, advance := newTimer(delay)
    defer stop()
    advance() // only has an effect when testing
    select {
    case <-ch:
        // We can proceed.
        return nil
    case <-ctx.Done():
        // Context was canceled before we could proceed.  Cancel the
        // reservation, which may permit other events to proceed sooner.
        r.Cancel()
        return ctx.Err()
    }

问题解析

经过主要函数分析，想必对limiter的实现，都有了足够的了解。然而，此处难免会产生一个新问题，limiter中的last，和lastEvent，分别代表着什么？调用limiter的ReserveN和AllowN时，允许传入时间参数t，在Reserve和Allow函数调用时，对应的时间参数t为当前时间。而在reserveN函数中，更新limiter时，t即为limiter的last，而lastEvent，代表着token已经被消耗的时间点，或者，换句话说，代表lastEvent之前的时间，已经没有更多的token可用了。针对这个结论，我们再代入到reservation的cancel方法验证。计算可返还的token时，取的时间差为lastEvent与reservation的timeToAct。在最终更新limiter状态时，last记录了取消时间，lastEvent则会根据返还的token，对应生成的时间，对比lastEvent。在上述的过程中，limiter的last不涉及token的计算，只是在记录着wait、cancel、allow、reserve发生的时间戳。而lastEvent，则与token的计算紧密相关。

总结

官方提供的limiter限流，提供了丰富的限流能力，同时，实现方式也保持着golang简单易懂的方式。在reserveN方法中，可以看出，有一些关于maxFutureReserve的判断，可以提前判断返回，降低mu锁持有的时间，提高limiter的并发性能，allow、allowN方法，是限流最普遍的用法，对应的maxFutureReserve均为0。