本地限流的go语言实现：官方包与uber包

前文中简介了几种典型的限流算法。本文对go语言中官方包golang.org/x/time/rate和uber包go.uber.org/ratelimit进行源码分析、性能对比、使用建议。

一、核心结论

官方包使用类GCRA的令牌桶变体算法，基于mutex保证并发安全，提供更丰富的功能。uber包使用GCRA算法的虚拟调度模式，基于atomic保证并发安全，提供极致的性能和极简的应用方式。作者更推荐大家在生产上使用官方包，因为对比之下，官方包的性能完全够用、功能更丰富、限速特性更容易理解。

二、性能对比

无并发竞争时，uber包是官方包性能的2.9倍，有并发竞争时，uber包性能是官方包的2.1倍。同时并发竞争均会降低限速器的性能，官方包下降80.81%，uber包下降85.94%。

看上去uber包性能更好，不过官方包在20 goroutinue的并发竞争下依然达到了167.6万qps，作者认为生产环境下完全够用了。

cpu: Intel(R) Xeon(R) CPU E5-2666 v3 @ 2.90GHz
BenchmarkOfficial-20            	 8732504	       121.0 ns/op	       0 B/op	       0 allocs/op
BenchmarkUber-20                	25541664	        47.32 ns/op	       0 B/op	       0 allocs/op
BenchmarkOfficialParallel-20    	 1676198	       731.1 ns/op	      51 B/op	       0 allocs/op
BenchmarkUberParallel-20        	 3590779	       326.6 ns/op	       0 B/op	       0 allocs/op

相关Banchmark代码

package main

import (
        "context"
        "testing"

        "go.uber.org/ratelimit"
        "golang.org/x/time/rate"
)

func BenchmarkOfficial(b *testing.B) {
        limit := rate.NewLimiter(100000000, 1)
        for i := 0; i < b.N; i++ {
                limit.Wait(context.Background())
        }
}

func BenchmarkUber(b *testing.B) {
        limit := ratelimit.New(100000000, ratelimit.WithSlack(1))
        for i := 0; i < b.N; i++ {
                limit.Take()
        }
}

func BenchmarkOfficialParallel(b *testing.B) {
        limit := rate.NewLimiter(100000000, 1)
        b.RunParallel(func(p *testing.PB) {
                for p.Next() {
                        limit.Wait(context.Background())
                }
        })
}

func BenchmarkUberParallel(b *testing.B) {
        limit := ratelimit.New(100000000, ratelimit.WithSlack(1))
        b.RunParallel(func(p *testing.PB) {
                for p.Next() {
                        limit.Take()
                }
        })
}

三、功能对比

标题	官方包	uber包
一次申请数量	1或N个	1个
突发应对	稳定	不稳定
超时控制	支持	不支持
阻塞请求	Wait/WaitN	Take
非阻塞请求	Allow/AllowN/Reserve/ReserveN	不支持
取消请求	Cancel	不支持

用法示例

func uber() {
	// 定义一个qps=10，burst能力=2的uber限速器
	limit := ratelimit.New(10, ratelimit.WithSlack(2))
	process := func() { log.Println("processed") }

	// 进行阻塞等待
	limit.Take()
	process()
}

func official() {
	// 定义一个qps=10，burst能力=2的官方限速器
	limit := rate.NewLimiter(10, 2)
	process := func() { log.Println("processed") }

	// 进行阻塞等待，并通过context配置等待超时
	ctx := context.Background()
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()
	err := limit.Wait(ctx)
	if err != nil {
		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
			// context超时会返回error
			// ...
		}
		// ...
		return
	}
	process()

	// 进行非阻塞的限流判断，可以对超出限速的流量进行抛弃、转异步等处理
	if limit.Allow() {
		process()
	} else {
		// 抛弃、转异步等
	}

	// Allow和Wait的低一级实现，应用方式更自由
	reserve := limit.Reserve()
	if reserve.OK() {
		process()
	} else {
		// 可以根据reserve.Delay()自行阻塞等待
		// 也可以调用reserve.Cancel()取消请求
		// 也可以抛弃、转异步等
	}
}

uber包突发应对能力不稳定问题

限流算法中的burst参数（uber实现中叫slack），对应生产中避免流量过于严格地塑形、让系统具有一定的突发应对能力的诉求。如平均qps=10时，10个请求可能集中在0.1秒到达，或者上一秒请求5个时当前秒是否允许放行15个请求。

官方算法可以理解为一个初始状态为满的令牌桶，因此初始的burst个请求能够马上取得令牌而不受限制。uber的burst能力却难以通过桶的比喻去理解，限速结果会让人感觉不稳定，这一点对生产不太友好------当你向同事推荐这个包时，你如何解释他的burst特性呢，难道要抛出下面这一坨case吗？

qps=1 busrt能力为3的官方限速器，可以理解为一个容量为3的桶、且初始为满令牌状态。
因此前3个请求能否马上取得令牌、不被限速，从第4个开始限速。
如果中途等待2秒，桶内又会积攒新的令牌，后续部分请求又能马上取得令牌。
official, do NO.1 at 0.0 second
official, do NO.2 at 0.0 second
official, do NO.3 at 0.0 second
official, do NO.4 at 1.0 second
official, do NO.5 at 2.0 second
等待2秒
official, do NO.6 at 4.0 second
official, do NO.7 at 4.0 second
official, do NO.8 at 5.0 second

但是同样定义一个burst能力为3的uber限速器，其特性并不稳定：
初始化限速器后马上开始请求，发现所有请求都被限速了，burst能力没有被体现
uber1, do 1 at 0.0 second
uber1, do 2 at 1.0 second
uber1, do 3 at 2.0 second
uber1, do 4 at 3.0 second
uber1, do 5 at 4.0 second
如果初始化限速器后等待3秒再开始请求，burst能力依然没有被体现
--等待3秒
uber2, do 1 at 0.0 second
uber2, do 2 at 1.0 second
uber2, do 3 at 2.0 second
uber2, do 4 at 3.0 second
uber2, do 5 at 4.0 second
只有在至少请求一次后，burst能力才会被累积
uber3, do 1 at 0.0 second
--等待3秒
uber3, do 2 at 3.0 second
uber3, do 3 at 3.0 second
uber3, do 4 at 3.0 second
uber3, do 5 at 4.0 second

测试代码

func main() {
	log.SetFlags(0)
	// 基于官方包定义一个1qps、burst能力为3的限速器
	log.Printf("qps=1 busrt能力为3的官方限速器，可以理解为一个容量为3的桶、且初始为满令牌状态。\n因此前3个请求能否马上取得令牌、不被限速，从第4个开始限速。\n如果中途等待2秒，桶内又会积攒新的令牌，后续部分请求又能马上取得令牌。")
	official := rate.NewLimiter(1, 3)
	start := time.Now()
	for i := 0; i < 8; i++ {
		err := official.Wait(context.Background())
		if err != nil {
			log.Fatal(err)
		}
		log.Printf("official, do NO.%d at %.1f second", i+1, time.Since(start).Seconds())
		if i == 4 {
			<-time.After(time.Second * 2)
			log.Println("等待2秒")
		}
	}

	log.Printf("\n但是同样定义一个burst能力为3的uber限速器，其特性并不稳定：")
	// 基于uber包定义一个1qps、burst能力为3的限速器
	log.Printf("初始化限速器后马上开始请求，发现所有请求都被限速了，burst能力没有被体现")
	uber1 := ratelimit.New(1, ratelimit.WithSlack(3))
	start = time.Now()
	for i := 0; i < 5; i++ {
		uber1.Take()
		log.Printf("uber1, do %d at %.1f second", i+1, time.Since(start).Seconds())
	}

	// 同上，但是在定义限速器后等待3秒
	log.Printf("如果初始化限速器后等待3秒再开始请求，burst能力依然没有被体现")
	uber2 := ratelimit.New(1, ratelimit.WithSlack(3))
	log.Printf("--等待3秒")
	<-time.After(3 * time.Second)
	start = time.Now()
	for i := 0; i < 5; i++ {
		uber2.Take()
		log.Printf("uber2, do %d at %.1f second", i+1, time.Since(start).Seconds())
	}

	// 同上，但是在第一次请求后等待3秒
	log.Printf("只有在至少请求一次后，burst能力才会被累积")
	uber3 := ratelimit.New(1, ratelimit.WithSlack(3))
	start = time.Now()
	for i := 0; i < 5; i++ {
		uber3.Take()
		log.Printf("uber3, do %d at %.1f second", i+1, time.Since(start).Seconds())
		if i == 0 {
			log.Println("--等待3秒")
			<-time.After(3 * time.Second)
		}
	}
}

四、源码分析

uber限流包核心代码

uber使用GCRA算法中的虚拟调度模式，根据上次请求到达时间、配置的速率和burst（slack），计算预期中的本次请求达到时间，再和当前真实时间进行对比、进行限速。过程中使用atomic包进行了并发安全的实现。

func (t *atomicInt64Limiter) Take() time.Time {
    var (
        // 本次请求的理论到达时间
        newTimeOfNextPermissionIssue int64
        // 当前时间
        now int64
    )
    for {
        now = t.clock.Now().UnixNano()
        // 上次请求到达时间
        timeOfNextPermissionIssue := atomic.LoadInt64(&t.state)
        switch {
        case timeOfNextPermissionIssue == 0 :
            // 首次请求：理论到达时间取当前时间
            newTimeOfNextPermissionIssue = now
        case (t.maxSlack == 0 && now-timeOfNextPermissionIssue > int64(t.perRequest)):
            // 可以放行本次请求，但不允许burst：理论到达时间取当前时间
            // perRequest：两次请求的最小间隔时间，second/qps
            // maxSlack：burst所需的时间，burst * perRequest
            // now-timeOfNextPermissionIssue：距离上次请求过了多久
            newTimeOfNextPermissionIssue = now
        case t.maxSlack > 0 && now-timeOfNextPermissionIssue > int64(t.maxSlack)+int64(t.perRequest):
            // 距离上一次请求，已经超出了burst所需的耗时：
            // 理论到达时间=当前时间-burst耗时，即最多保留burst耗时的冗余
            newTimeOfNextPermissionIssue = now - int64(t.maxSlack)
        default:
            // 其余情况：理论到达时间为上次到达时间+请求间隔
            newTimeOfNextPermissionIssue = timeOfNextPermissionIssue + int64(t.perRequest)
        }
        // 进行并发竞争
        if atomic.CompareAndSwapInt64(&t.state, timeOfNextPermissionIssue, newTimeOfNextPermissionIssue) {
            // 竞争成功
            break
        }
        // 竞争失败，需要重新计算理论到达时间
    }
    // 等待到理论到达时间，放行本次请求
    sleepDuration := time.Duration(newTimeOfNextPermissionIssue - now)
    if sleepDuration > 0 {
        t.clock.Sleep(sleepDuration)
        return time.Unix(0, newTimeOfNextPermissionIssue)
    }
    return time.Unix(0, now)
}

官方限流包核心代码

官方包生成自己实现的是令牌桶算法，并给出了wiki的相关链接，如果按令牌桶去理解的话，官方包实现的桶特性是：初始满状态、桶容量=burst、按预期rate向桶中添加令牌、令牌满了无法再添加、请求到达时需要从桶中拿出一个令牌才放行。

不过具体实现上，官方包的做法更接近于GCRA算法的连续桶模式，请求到达时根据上次请求时的桶状态、当前时间来计算最新的桶状态，再进行限流判断。另外官方包通过mutex进行并发安全的实现。

// Wait 阻塞等待，直至成功请求到一个令牌
func (lim *Limiter) Wait(ctx context.Context) (err error) {
    return lim.WaitN(ctx, 1)
}
// WaitN 阻塞等待，直至成功请求到n个令牌
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
    // newTimer：用于等待请求的定时器，默认基于time.Timer实现
    newTimer := ...
    return lim.wait(ctx, n, time.Now(), newTimer)
}
// wait 阻塞等待的实现
func (lim *Limiter) wait(ctx context.Context, n int, t time.Time, newTimer func(d time.Duration) (<-chan time.Time, func() bool, func())) error {
    
    ...
    // 请求令牌数不能超出桶大小
    if n > burst && limit != Inf {
        return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst)
    }
    // 结合ctx中的超时设置，确定最大等待时间
    waitLimit := ...
    // 尝试获取n个令牌
    r := lim.reserveN(t, n, waitLimit)
    // 获取失败
    if !r.ok {
        return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
    }
    // 获取成功，但如果delay>0则说明是预支的令牌、需要阻塞等待直到这些令牌被生成
    delay := r.DelayFrom(t)
    if delay == 0 {
        return nil
    }
    
    // 生成定时器、执行阻塞等待
    ...
}

// reserveN 令牌桶的核心实现，尝试在t时间扣除n个令牌
func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration) Reservation {
    // 依赖于锁
    lim.mu.Lock()
    defer lim.mu.Unlock()
    // 特殊case处理
    ...
    // 计算在t时，桶内令牌的数量tokens
    t, tokens := lim.advance(t)
    // 尝试扣除n个令牌
    tokens -= float64(n)
    // 令牌不足时，计算生成足够令牌所需的等待时间
    var waitDuration time.Duration
    if tokens < 0 {
        waitDuration = lim.limit.durationFromTokens(-tokens)
    }
    // 放行条件：
    // n未超过burst能力，即同时取的令牌数未超出桶容量
    // 且所需等待时间未超出调用者设置的最大等待时间
    ok := n <= lim.burst && waitDuration <= maxFutureReserve
    // 生成令牌扣除结果
    r := ...
    if ok {
        ...
        // 更新桶快照
        lim.last = t                // 桶状态更新时间
        lim.tokens = tokens         // 桶内令牌数
        lim.lastEvent = r.timeToAct // 本次n个请求的最早执行完毕时间，用于退回令牌本次获取的令牌Cancel
    }
    return r
}
// advance 计算到t时，桶内的令牌数
func (lim *Limiter) advance(t time.Time) (newT time.Time, newTokens float64) {
    // 桶状态的上一次快照时间
    last := lim.last
    // 兼容过期请求
    if t.Before(last) {
        last = t
    }
    // 根据上一次快照时间、本次事件，计算生成的令牌数
    elapsed := t.Sub(last)
    delta := lim.limit.tokensFromDuration(elapsed)
    // 更新桶内令牌数，且不能超出桶容量（burst能力）
    tokens := lim.tokens + delta
    if burst := float64(lim.burst); tokens > burst {
        tokens = burst
    }
    // 返回本次桶快照时间、桶内令牌数
    return t, tokens
}