首发于技术拷问

[golang]rate 限制器

1.背景

golang.org/x/time/rate 限制器提供了功能来控制事件发生的频率,在项目中遇到。研究一下。

对应文档 godoc.org/golang.org/x/

2.核心结构

Limiter

Limiter对象: 控制事件发生的频率。它实现了一个令牌桶。开始的时候为满的,大小为b。然后每秒补充r个令牌。如果r取Inf(无穷大),则忽略b。

Limiter的默认初始化(Zero Value)是一个有效值,但是会拒绝所有的事件。需要使用NewLimiter来创建实际可用的限速器。

Limiter有三个主要方法Allow,Reserve和Wait,大多数情况都应该使用Wait。每个方法的调用都会消费一个或N个(针对WaitN之类)单独的token,在没有toekn可用时三者的表现不同。

  • Allow 无可用token则返回false
  • Reserve 无可用token,则返回一个或多个未来token的预订以及调用者在使用前必须等待的时长。
  • Wait无可用token会阻塞住,直到获取一个token,或者超时或取消(基于context.Context

Reservation

Reservation 保存的是指定delay时间段后Limiter允许的事件信息。一个Reservation可能被取消掉,这样Limiter就允许更多的事件被其他调用方处理。

3.核心代码

// 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.
// 输入预约到的时间now, 预约的数量n,以及最大可接受的等待时间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,
        }
    }

    now, last, tokens := lim.advance(now) //这个函数返回指定到预约时间后当时Limiter会处于的状态。last表示其数据对应的最新时间;tokens表示到预约的时间点上能有多少令牌存在。

    // Calculate the remaining number of tokens resulting from the request.
    tokens -= float64(n) //去掉预约的数量后,Limiter上还会剩余多少令牌。

    // Calculate the wait duration
    var waitDuration time.Duration
    if tokens < 0 {
        waitDuration = lim.limit.durationFromTokens(-tokens)
    }

    // Decide result
    ok := n <= lim.burst && waitDuration <= maxFutureReserve

    // Prepare reservation
    r := Reservation{
        ok:    ok,
        lim:   lim,
        limit: lim.limit,
    }
    if ok {
        r.tokens = n
        r.timeToAct = now.Add(waitDuration)
    }

    // Update state
    if ok {
        lim.last = now          //token数据最后更新的时间
        lim.tokens = tokens     //这里可能会是负数
        lim.lastEvent = r.timeToAct     //限制器事件发生的最近一次,可能是以前或者未来
    } else {
        lim.last = last 
    }

    lim.mu.Unlock()
    return r
}

代码非常简洁,每次调用才会进行计算,非常高效。而且有加锁机制,线程安全。可以很放心的使用!

4.使用示例

Wait

package main

import (
    "golang.org/x/time/rate"
    "fmt"
    "context"
    "time"
    "sync"
)

func wait() {
    limiter := rate.NewLimiter(3, 5)
    ctx, cancel := context.WithTimeout(context.Background(), time.Second * 5)
    defer cancel()
    for i:=0; ; i++ {
        fmt.Printf("%03d %s\n", i, time.Now().Format("2006-01-02 15:04:05.000"))
        err := limiter.Wait(ctx)
        if err != nil {
            fmt.Printf("err: %s\n", err.Error())
            return
        }
    }
}

func main() {
    allow()
}

这里指定令牌桶大小为5,每秒补充3个令牌。这个算法的逻辑是有多少都直接消耗,速度极快。所以第一秒内理论上能消耗初始化的5以及该秒内补充的3,后续4秒每秒消耗完补充的3。合计5+3*5(超时时间为5秒)=20,从输出上看即符合预期。

输出信息如下:

000 2018-11-21 20:00:08.304
001 2018-11-21 20:00:08.304
002 2018-11-21 20:00:08.304
003 2018-11-21 20:00:08.304
004 2018-11-21 20:00:08.304
005 2018-11-21 20:00:08.304
006 2018-11-21 20:00:08.642
007 2018-11-21 20:00:08.971
008 2018-11-21 20:00:09.304
009 2018-11-21 20:00:09.639
010 2018-11-21 20:00:09.971
011 2018-11-21 20:00:10.304
012 2018-11-21 20:00:10.638
013 2018-11-21 20:00:10.972
014 2018-11-21 20:00:11.304
015 2018-11-21 20:00:11.639
016 2018-11-21 20:00:11.972
017 2018-11-21 20:00:12.305
018 2018-11-21 20:00:12.638
019 2018-11-21 20:00:12.975
err: rate: Wait(n=1) would exceed context deadline

Reserve

func reserve() {
    limiter := rate.NewLimiter(3, 5)
    ctx, cancel := context.WithTimeout(context.Background(), time.Second * 5)
    defer cancel()
    for i:=0; ; i++ {
        fmt.Printf("%03d %s\n", i, time.Now().Format("2006-01-02 15:04:05.000"))
        reserve := limiter.Reserve()
        if ! reserve.OK() {
            //返回是异常的,不能正常使用
            fmt.Println("Not allowed to act! Did you remember to set lim.burst to be > 0 ?")
            return
        }
        delayD := reserve.Delay()
        fmt.Println("sleep delay ", delayD)
        time.Sleep(delayD)
        select {
        case <- ctx.Done():
            fmt.Println("timeout, quit")
            return
        default:
        }
       //TODO 业务逻辑
    }
}

//本函数仅仅用于演示并发的安全性,实际使用如下代码来Sleep
//delayD := reserve.Delay()
//fmt.Println("sleep delay ", delayD)
//time.Sleep(delayD)
func reserve2() {
    limiter := rate.NewLimiter(3, 5)
    var wg = sync.WaitGroup{}
    for i:=0; i<5; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            r := limiter.ReserveN(time.Now(), 5)
            if r.OK() {
                fmt.Printf("r1 delay %s\n", r.Delay())
            } else {
                fmt.Println("not ok")
            }
        }()
    }

    wg.Wait()
}

Allow

func allow() {
    limiter := rate.NewLimiter(3, 5)
    for i:=0; i<50; i++ {
        if  limiter.Allow() {
            fmt.Printf("%03d Ok  %s\n", i, time.Now().Format("2006-01-02 15:04:05.000"))
        } else {
            fmt.Printf("%03d Err %s\n", i, time.Now().Format("2006-01-02 15:04:05.000"))
        }
        time.Sleep(100*time.Millisecond)
    }
}

很好用的一个包,有待进一步发掘功能与场景,Good luck!

参考

jianshu.com/p/4ce68a31a

en.wikipedia.org/wiki/T

发布于 2020-01-01

文章被以下专栏收录