目录
前言例子实现小结

前言

在翻Golang官方库的过程中,发现一个有趣的库golang.org/x/time ,里面只有一个类rate,研究了一下发现它是一个限流器,实现了很多的功能,当然它的核心原理并不复杂,也就是令牌桶算法。

令牌桶算法的原理是:令牌桶会不断地把令牌添加到桶里,而请求会从桶中获取令牌,只有拥有令牌地请求才能被接受。因为桶中可以提前保留一些令牌,所以它允许一定地突发流量通过。

例子

下面是限流算法常见的写法,首先判断是否有令牌,如果有就通过,否则直接失败。


(资料图片)

package main
import (
	"fmt"
	"time"
	"golang.org/x/time/rate"
)
func main() {
        // 每0.1秒生成一个令牌,也就是一秒10个令牌,最大保留令牌上限10
	l := rate.NewLimiter(rate.Every(time.Second/10), 10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 判断是否有令牌,如果有就输出
				if l.Allow() {
					fmt.Printf("allow %d\n", i)
				}
                                // 每0.5秒请求一次
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

上面的rate.Every(time.Second/10)会返回一个Limit类型,代表每秒生成多少个令牌。

这个库还提供了另外一种写法,等待直到有令牌为止(或超时):

func main() {
	l := rate.NewLimiter(rate.Every(time.Second/10), 100)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 等待直到有令牌
				if err := l.Wait(context.TODO()); err != nil {
				} else {
					fmt.Printf("allow %d\n", i)
				}
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

这样在某些场景下我们可以让请求等待一会,而不是直接失败。

还有一个更加特殊的请求令牌方式,也就是先预留令牌,到指定时间不再需要去获取令牌,直接执行操作即可:

func main() {
	l := rate.NewLimiter(rate.Every(time.Second/10), 10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			for {
                                // 先预留令牌
				if r := l.Reserve(); r.OK() {
                                        // 休眠直到令牌生效
					time.Sleep(r.Delay())
					fmt.Printf("allow %d\n", i)
				}
				time.Sleep(time.Second / 2)
			}
		}(i)
	}
	time.Sleep(time.Second * 10)
}

当然,如果预留的令牌不想使用了,也可以使用r.Cancel()归还已预留的令牌。

上面的Allow()、Wait()、Reserve()都是一次消耗一个令牌,其实都有对应的AllowN()、WaitN()、ReserveN()方法,一次消耗N个令牌,这样就可以根据任务消耗的资源灵活的消耗令牌。

实现

不管我们从Allow()、Wait()还是Reserve()进去,最终都会进入到reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation 方法:

首先在进入方法的时候,会先处理两种特殊情况:

// 如果无限生成令牌,则直接返回
        if lim.limit == Inf {
		return Reservation{
			ok:        true,
			lim:       lim,
			tokens:    n,
			timeToAct: now,
		}
        // 如果不会生成令牌,则在初始令牌里面拿,直到拿完为止
	} else if lim.limit == 0 {
		var ok bool
		if lim.burst >= n {
			ok = true
			lim.burst -= n
		}
		return Reservation{
			ok:        ok,
			lim:       lim,
			tokens:    lim.burst,
			timeToAct: now,
		}
	}

然后重新计算当前有多少令牌,减去要消耗的令牌:

// 计算当前有多少令牌
	now, last, tokens := lim.advance(now)
	// 减去要消耗的N个令牌
	tokens -= float64(n)
        // 如果剩余令牌为负数,那么计算一下要等待多久才能拿到令牌
        var waitDuration time.Duration
	if tokens < 0 {
		waitDuration = lim.limit.durationFromTokens(-tokens)
	}
        // 判断请求是否成功,maxFutureReserve代表最大可以等待的时间,也就是请求能否接收拿到令牌需要等待的时间
        ok := n <= lim.burst && waitDuration <= maxFutureReserve

余下的代码就是更新限流器的状态。

小结

可以看到这个令牌桶限流器实现的功能非常的丰富,如果需要令牌桶限流器,可以优先考虑使用这个实现。

关键词: 还有一个 特殊情况 可以使用 直接执行