golang 令牌桶 rate

2024-10-16 19:39:02 编程来源：ZONE.CI 全球网 0 阅读模式

令牌桶算法（Token Bucket）是一种常用的限流算法，广泛应用于互联网系统中。在 Go 语言中，我们可以通过使用 golang.org/x/time/rate 包来实现令牌桶算法。Rate 结构体提供了强大的功能，可以用于控制函数或服务的请求速率。接下来，我将详细介绍 golang 令牌桶 Rate 包的使用。

使用 Rate 限制请求速率

在开发中，我们经常需要控制对其他系统的请求速率，避免过多的请求导致系统负载过高或者被其他系统封禁。golang.org/x/time/rate 包提供了 Rate 结构体，可以非常方便地控制请求速率。

首先，我们需要创建一个 Rate 实例：

limiter := rate.NewLimiter(rate.Limit(10), 3)

上述代码表示我们希望每秒最多处理 10 个请求，并且系统启动后的前 3 秒允许突发处理更多的请求。Limiter 的第一个参数是 Limit 类型，可以使用 time.Second、time.Minute 等单位来表示时间间隔。

限制请求速率

Rate 实例创建完成后，我们可以使用 Wait 方法来等待令牌，以控制请求速率：

err := limiter.Wait(context.TODO())
if err != nil {
    fmt.Println("请求被限制")
} else {
    fmt.Println("处理请求")
}

当使用 Wait 方法时，如果当前时间可以立即获取到令牌，则 Wait 方法会立即返回。如果当前时间无法立即获取到令牌，则 Wait 方法会阻塞，直到能够获取到令牌或者后续操作超时。

限制并发请求

有时候，我们需要限制同时发起的请求数量，以保证系统资源的合理分配。golang.org/x/time/rate 结合 golang.org/x/sync/semaphore 包可以很容易地实现此类限制。

首先，我们需要创建一个 Rate 实例和一个 Semaphore（信号量）实例：

limiter := rate.NewLimiter(rate.Limit(10), 3)
sem := semaphore.NewWeighted(5)

上述代码表示我们希望每秒最多处理 10 个请求，并且系统启动后的前 3 秒允许突发处理更多的请求。Semaphore 的参数表示可以同时处理的请求数量。

在实际处理请求前，我们可以使用 Acquire 方法来申请一个信号量：

ctx := context.TODO()
err := sem.Acquire(ctx, 1)
if err != nil {
    fmt.Println("请求被限制")
    return
}
defer sem.Release(1)

err = limiter.Wait(ctx)
if err != nil {
    fmt.Println("请求被限制")
    return
}

fmt.Println("处理请求")

在上述代码中，首先我们使用 Acquire 方法申请一个信号量，如果能够获得信号量，则继续下一步处理请求；否则，说明当前请求数已达到限制，无法继续处理。

在处理请求完毕后，我们使用 Release 方法释放一个信号量。