Golang官方限流器库实现限流示例详解

尹扇殞

前言

在翻Golang官方库的过程中，发现一个有趣的库golang.org/x/time ，里面只有一个类rate，研究了一下发现它是一个限流器，实现了很多的功能，当然它的核心原理并不复杂，也就是令牌桶算法。
令牌桶算法的原理是：令牌桶会不断地把令牌添加到桶里，而请求会从桶中获取令牌，只有拥有令牌地请求才能被接受。因为桶中可以提前保留一些令牌，所以它允许一定地突发流量通过。

例子

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

1. package main
   
2. import (
   
3.         "fmt"
   
4.         "time"
   
5.         "golang.org/x/time/rate"
   
6. )
   
7. func main() {
   
8.         // 每0.1秒生成一个令牌，也就是一秒10个令牌，最大保留令牌上限10
   
9.         l := rate.NewLimiter(rate.Every(time.Second/10), 10)
   
10.         for i := 0; i < 10; i++ {
    
11.                 go func(i int) {
    
12.                         for {
    
13.                                 // 判断是否有令牌，如果有就输出
    
14.                                 if l.Allow() {
    
15.                                         fmt.Printf("allow %d\n", i)
    
16.                                 }
    
17.                                 // 每0.5秒请求一次
    
18.                                 time.Sleep(time.Second / 2)
    
19.                         }
    
20.                 }(i)
    
21.         }
    
22.         time.Sleep(time.Second * 10)
    
23. }

复制代码

上面的rate.Every(time.Second/10)会返回一个Limit类型，代表每秒生成多少个令牌。
这个库还提供了另外一种写法，等待直到有令牌为止（或超时）：

1. func main() {
   
2.         l := rate.NewLimiter(rate.Every(time.Second/10), 100)
   
3.         for i := 0; i < 10; i++ {
   
4.                 go func(i int) {
   
5.                         for {
   
6.                                 // 等待直到有令牌
   
7.                                 if err := l.Wait(context.TODO()); err != nil {
   
8.                                 } else {
   
9.                                         fmt.Printf("allow %d\n", i)
   
10.                                 }
    
11.                                 time.Sleep(time.Second / 2)
    
12.                         }
    
13.                 }(i)
    
14.         }
    
15.         time.Sleep(time.Second * 10)
    
16. }

复制代码

这样在某些场景下我们可以让请求等待一会，而不是直接失败。
还有一个更加特殊的请求令牌方式，也就是先预留令牌，到指定时间不再需要去获取令牌，直接执行操作即可：

1. func main() {
   
2.         l := rate.NewLimiter(rate.Every(time.Second/10), 10)
   
3.         for i := 0; i < 10; i++ {
   
4.                 go func(i int) {
   
5.                         for {
   
6.                                 // 先预留令牌
   
7.                                 if r := l.Reserve(); r.OK() {
   
8.                                         // 休眠直到令牌生效
   
9.                                         time.Sleep(r.Delay())
   
10.                                         fmt.Printf("allow %d\n", i)
    
11.                                 }
    
12.                                 time.Sleep(time.Second / 2)
    
13.                         }
    
14.                 }(i)
    
15.         }
    
16.         time.Sleep(time.Second * 10)
    
17. }

复制代码

当然，如果预留的令牌不想使用了，也可以使用r.Cancel()归还已预留的令牌。
上面的Allow()、Wait()、Reserve()都是一次消耗一个令牌，其实都有对应的AllowN()、WaitN()、ReserveN()方法，一次消耗N个令牌，这样就可以根据任务消耗的资源灵活的消耗令牌。

实现

不管我们从Allow()、Wait()还是Reserve()进去，最终都会进入到reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation 方法：
首先在进入方法的时候，会先处理两种特殊情况：

1.         // 如果无限生成令牌，则直接返回
   
2.         if lim.limit == Inf {
   
3.                 return Reservation{
   
4.                         ok:        true,
   
5.                         lim:       lim,
   
6.                         tokens:    n,
   
7.                         timeToAct: now,
   
8.                 }
   
9.         // 如果不会生成令牌，则在初始令牌里面拿，直到拿完为止
   
10.         } else if lim.limit == 0 {
    
11.                 var ok bool
    
12.                 if lim.burst >= n {
    
13.                         ok = true
    
14.                         lim.burst -= n
    
15.                 }
    
16.                 return Reservation{
    
17.                         ok:        ok,
    
18.                         lim:       lim,
    
19.                         tokens:    lim.burst,
    
20.                         timeToAct: now,
    
21.                 }
    
22.         }

复制代码

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

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

复制代码

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

小结

可以看到这个令牌桶限流器实现的功能非常的丰富，如果需要令牌桶限流器，可以优先考虑使用这个实现。
以上就是Golang官方限流器库使用示例详解的详细内容，更多关于Golang官方限流器库的资料请关注晓枫资讯其它相关文章！

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！

颜昭昭

顶顶更健康！！！

克罗地亚之夜之o

路过，支持一下