在研发中，我们经常会面对到处理并发逻辑的场景，尤其是有时候在与第三方平台对接的场景下，会遇到请求限流限QPS的要求。对于限流或者限速，我们通常会用两种算法来满足需要：
- 令牌桶算法：在特定容量的桶里面装令牌，当令牌数量小于桶的容量时，会持续以我们预期的限流速率生产令牌；不管桶里面是不是空的，业务都得等到拿到令牌，才能继续执行业务逻辑
- 漏桶算法：业务先统一进入桶里，桶满了之后，会以我们预期的限流速率，一个个把在等待的业务漏出去，然后各个业务才开始执行业务逻辑
这两种算法，虽然实际QPS数值可能都会有波动，但都能把速率限制在一个合理的水位。在Golang里面，这两种算法都有现成的实现可以直接用。咱们今天，就来看看这块的例子。
首先来看看令牌桶算法，在Golang自带的rate库中，就有了一份令牌桶算法的实现。我们上一下代码的例子，来看一下rate库的基本用法：
import "golang.org/x/time/rate" // 需要import的rate库，其它import暂时忽略// 生成0->X的数据集
func generateData(num int) []any {var data []anyfor i := 0; i < num; i++ {data = append(data, i)}return data
}// 处理数据，数字*10
func process(obj any) (any, error) {integer, ok := obj.(int)if !ok {return nil, errors.New("invalid integer")}time.Sleep(1)nextInteger := integer * 10if integer%99 == 0 {return nextInteger, errors.New("not a happy number")}return nextInteger, nil
}func TestRate(t *testing.T) {limit := rate.Limit(50) // QPS：50burst := 25 // 桶容量25limiter := rate.NewLimiter(limit, burst)size := 500 // 数据量500data := generateData(size)var wg sync.WaitGroupstartTime := time.Now()for i, item := range data {wg.Add(1)go func(idx int, obj any) {defer wg.Done()// 拿到令牌if err := limiter.Wait(context.Background()); err != nil {t.Logf("[%d] [EXCEPTION] wait err: %v", idx, err)}// 执行业务逻辑processed, err := process(obj)if err != nil {t.Logf("[%d] [ERROR] processed: %v, err: %v", idx, processed, err)} else {t.Logf("[%d] [OK] processed: %v", idx, processed)}}(i, item)}wg.Wait()endTime := time.Now()t.Logf("start: %v, end: %v, seconds: %v", startTime, endTime, endTime.Sub(startTime).Seconds())
}
通过limiter := rate.NewLimiter(limit, burst)的形式，我们可以初始化一个令牌生成速率为limit，容量为burst的令牌桶。在业务里，则通过limiter.Wait(ctx)的方式拿到一个令牌，执行逻辑。结合先前讲过的WaitGroup，我们就能够实现开多个goroutine异步执行任务，并配合limiter来做业务逻辑的限速。
这里需要注意，从长线来看，limit的大小是能够决定一个基础的限速速率，但从短线角度来看，burst这个桶具备了缓冲作用，在冷启动时，由于burst的存在，初始的QPS会比实际预估的较大。因此，业务通过令牌桶方式限速时，如果需要限制冷启动时的瞬时速率，需要留意把burst的值设置小一些。
然后我们再来看通过漏桶算法来限速的例子，这回需要用到一个开源库uber-go/ratelimit来实现。例子也非常简单，代码如下：
func TestRateLimit(t *testing.T) {limiter := ratelimit.New(50) // 漏桶速率50size := 500data := generateData(size)var wg sync.WaitGroupstartTime := time.Now()for i, item := range data {wg.Add(1)go func(idx int, obj any) {defer wg.Done()limiter.Take() // 入桶待漏processed, err := process(obj)if err != nil {t.Logf("[%d] [ERROR] processed: %v, err: %v", idx, processed, err)} else {t.Logf("[%d] [OK] processed: %v", idx, processed)}}(i, item)}wg.Wait()endTime := time.Now()t.Logf("start: %v, end: %v, seconds: %v", startTime, endTime, endTime.Sub(startTime).Seconds())
}
ratelimit库速率的单位可以类比为QPS，通过ratelimit.New，可以初始化一个指定QPS限制的限流器。执行业务逻辑前，需要通过limiter.Take()逻辑等待漏出来后，才能执行后续的逻辑。这样便实现了限流的效果。
