golang 限流器 time/rate

golang.org/x/time/rate
基于令牌桶算法：随时间以 1/r 个令牌的速度向容积为 b 个令牌的桶中添加令牌，有请求就取走令牌，若令牌不足则不执行请求或者等待。

加令牌的时机

通过初始化函数可以看到，并没有开启一个定时任务定时添加令牌，所以猜测是在每次请求令牌之前，通过时间添加固定的令牌。

func NewLimiter(r Limit, b int) *Limiter {
    return &Limiter{
        limit: r,
        burst: b,
    }
}

其中，limit 表示限流器添加的速率，burst 表示总的容量

其主要是通过 advance 函数计算当前时间可以添加的令牌数，基本上每个方法都会执行这个函数

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
}

两个重要的方法

1. AllowN ：是否允许执行

package main

import (
    "fmt"
    "time"

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

func main() {
    r := rate.Every(1 * time.Second)
    limit := rate.NewLimiter(r, 10)
    for i := 0; i < 5; i++ {
        if ok := limit.AllowN(time.Now(), 3); ok {
            fmt.Println("next")
        } else {
            fmt.Println("limit")
            break
        }
    }
}

2. ReserveN：返回一个 Reservation ，指示调用者在 n 个事件发生之前必须等待多长时间。

package main

import (
    "fmt"
    "time"

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

func main() {
    r := rate.Every(1 * time.Second)
    limit := rate.NewLimiter(r, 10)
    for i := 0; i < 5; i++ {
        r := limit.ReserveN(time.Now(), 8)
        fmt.Println(r.Delay())
    }
}

输出：

0s
5.999962161s
13.999959821s
21.999958261s
29.999956713s

这两个方法都调用了 reserveN 方法，这个方法也是整个限流器里面最重要的方法之一。
简单来说就是：首先通过时间添加令牌，再检查令牌的个数，计算需要等待的时间。