【从零单排Golang】第十四话:使用rate和ratelimit实现限流限速

在研发中,我们经常会面对到处理并发逻辑的场景,尤其是有时候在与第三方平台对接的场景下,会遇到请求限流限QPS的要求。对于限流或者限速,我们通常会用两种算法来满足需要:

  • 令牌桶算法:在特定容量的桶里面装令牌,当令牌数量小于桶的容量时,会持续以我们预期的限流速率生产令牌;不管桶里面是不是空的,业务都得等到拿到令牌,才能继续执行业务逻辑
  • 漏桶算法:业务先统一进入桶里,桶满了之后,会以我们预期的限流速率,一个个把在等待的业务漏出去,然后各个业务才开始执行业务逻辑

这两种算法,虽然实际QPS数值可能都会有波动,但都能把速率限制在一个合理的水位。在Golang里面,这两种算法都有现成的实现可以直接用。咱们今天,就来看看这块的例子。

首先来看看令牌桶算法,在Golang自带的rate库中,就有了一份令牌桶算法的实现。我们上一下代码的例子,来看一下rate库的基本用法:

go 复制代码
import "golang.org/x/time/rate" // 需要import的rate库,其它import暂时忽略

// 生成0->X的数据集
func generateData(num int) []any {
    var data []any
    for 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 * 10
    if integer%99 == 0 {
        return nextInteger, errors.New("not a happy number")
    }
    return nextInteger, nil
}

func TestRate(t *testing.T) {
    limit := rate.Limit(50) // QPS:50
    burst := 25 // 桶容量25
    limiter := rate.NewLimiter(limit, burst)
    size := 500 // 数据量500

    data := generateData(size)
    var wg sync.WaitGroup
    startTime := 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来实现。例子也非常简单,代码如下:

go 复制代码
func TestRateLimit(t *testing.T) {
    limiter := ratelimit.New(50) // 漏桶速率50
    size := 500

    data := generateData(size)
    var wg sync.WaitGroup
    startTime := 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()逻辑等待漏出来后,才能执行后续的逻辑。这样便实现了限流的效果。

相关推荐
tyler_download1 小时前
golang 实现比特币内核:处理椭圆曲线中的天文数字
golang·blockchain·bitcoin
疯狂的程需猿2 小时前
一个百度、必应搜索引擎图片获取下载的工具包
golang·图搜索
明月看潮生4 小时前
青少年编程与数学 02-003 Go语言网络编程 09课题、Cookie
青少年编程·golang·网络编程·编程与数学
明月看潮生4 小时前
青少年编程与数学 02-003 Go语言网络编程 15课题、Go语言URL编程
开发语言·网络·青少年编程·golang·编程与数学
明月看潮生4 小时前
青少年编程与数学 02-003 Go语言网络编程 14课题、Go语言Udp编程
青少年编程·golang·网络编程·编程与数学
hlsd#6 小时前
go 集成go-redis 缓存操作
redis·缓存·golang
qq_1728055919 小时前
GIN 反向代理功能
后端·golang·go
__AtYou__20 小时前
Golang | Leetcode Golang题解之第535题TinyURL的加密与解密
leetcode·golang·题解
kevin_tech1 天前
Go API 多种响应的规范化处理和简化策略
开发语言·后端·golang·状态模式
幺零九零零1 天前
【Golang】sql.Null* 类型使用(处理空值和零值)
数据库·sql·golang