RateLimiter 是一个用来限制访问速率的工具,让你可以控制在一段时间内某个处理单元能够处理的最大请求次数。下面是一个简单的 Go 语言实现的 RateLimiter 示例代码:
package main
import (
"fmt"
"time"
)
type RateLimiter struct {
rate int // 限制速率,单位:次/秒
burst int // 突发频率,即可以瞬间处理的请求次数
batchSize int // 批次大小,即每次处理的最大请求次数
tokens int // 当前令牌数
lastUpdate time.Time // 上次更新令牌数时间
mu sync.Mutex // 互斥锁,用来实现 goroutine 安全
}
// 初始化 RateLimiter
func New(rate, burst, batchSize int) *RateLimiter {
return &RateLimiter{
rate: rate,
burst: burst,
batchSize: batchSize,
tokens: burst,
lastUpdate: time.Now(),
}
}
// 进行请求速率控制
func (r *RateLimiter) Wait() {
r.mu.Lock()
// 计算需要等待的时间
delay := (r.batchSize + r.tokens - r.burst) / r.rate
now := time.Now()
next := r.lastUpdate.Add(time.Duration(delay) * time.Second)
if now.Before(next) {
time.Sleep(next.Sub(now))
}
// 更新令牌数
elapsed := int(now.Sub(r.lastUpdate).Seconds())
r.tokens = min(r.burst, r.tokens+elapsed*r.rate)
r.lastUpdate = now
// 判断是否可以执行请求
if r.tokens < r.batchSize {
r.mu.Unlock()
r.Wait()
} else {
r.tokens -= r.batchSize
r.mu.Unlock()
}
}
// 取较小值函数
func min(a, b int) int {
if a < b {
return a
}
return b
}
func main() {
// 初始化 RateLimiter
rl := New(1, 5, 2)
// 并发发送 10 次请求
for i := 0; i < 10; i++ {
go func() {
rl.Wait()
fmt.Println("处理请求...")
}()
}
// 等待所有请求完成
time.Sleep(10 * time.Second)
}
在上面的示例代码中,我们定义了一个 RateLimiter 结构体,它包含了限制速率、突发频率、批次大小、当前令牌数和上次更新时间等属性。我们利用互斥锁来保证goroutine的安全。
Wait() 方法是进行请求速率控制的核心方法。它首先获取互斥锁并计算需要等待的时间,然后利用 time.Sleep() 等待这段时间,然后根据经过的时间更新令牌数。如果当前令牌数不足以处理请求,则递归等待。
最后,我们在 main() 函数中创建了 10 个 goroutine 并发发送请求,利用 time.Sleep() 等待它们完成后退出程序。