go的令牌桶实现
golang.org/x/time/rate
包是 Go 语言的一个扩展包,用于实现基于令牌桶算法的限流器。这个包非常适合用于控制对资源的访问速率,比如 API 请求的速率限制。
示例演示
我们首先通过一个简单的例子,来展示怎么实现限流的。在这个例子中,我们将创建一个限流器,它允许每秒最多通过 2个请求,允许突发6个请求,我们将模拟多个请求来观察限流效果。
package main
import (
"fmt"
"golang.org/x/time/rate"
"time"
)
func main() {
// 创建一个每秒允许2个令牌的Limiter,桶的容量为6
limiter := rate.NewLimiter(2, 6)
for i := 0; i < 100; i++ {
if limiter.Allow() {
fmt.Println("Request", i, "is allowed at", time.Now().Format("15:04:05"))
} else {
fmt.Println("Request", i, "is not allowed at", time.Now().Format("15:04:05"))
}
// 模拟请求间隔
time.Sleep(100 * time.Millisecond)
}
}
输出如下:
Request 0 is allowed at 16:14:28
Request 1 is allowed at 16:14:28
Request 2 is allowed at 16:14:28
Request 3 is allowed at 16:14:28
Request 4 is allowed at 16:14:28
Request 5 is allowed at 16:14:28
Request 6 is allowed at 16:14:29
Request 7 is not allowed at 16:14:29
Request 8 is not allowed at 16:14:29
Request 9 is not allowed at 16:14:29
Request 10 is allowed at 16:14:29
Request 11 is not allowed at 16:14:29
Request 12 is not allowed at 16:14:29
Request 13 is not allowed at 16:14:29
Request 14 is allowed at 16:14:29
Request 15 is not allowed at 16:14:30
Request 16 is not allowed at 16:14:30
Request 17 is not allowed at 16:14:30
Request 18 is not allowed at 16:14:30
Request 19 is allowed at 16:14:30
Request 20 is not allowed at 16:14:30
Request 21 is not allowed at 16:14:30
Request 22 is not allowed at 16:14:30
Request 23 is allowed at 16:14:30
Request 24 is not allowed at 16:14:31
Request 25 is not allowed at 16:14:31
Request 26 is not allowed at 16:14:31
Request 27 is not allowed at 16:14:31
Request 28 is allowed at 16:14:31
Request 29 is not allowed at 16:14:31
Request 30 is not allowed at 16:14:31
Request 31 is not allowed at 16:14:31
Request 32 is not allowed at 16:14:31
Request 33 is allowed at 16:14:32
Request 34 is not allowed at 16:14:32
Request 35 is not allowed at 16:14:32
Request 36 is not allowed at 16:14:32
Request 37 is allowed at 16:14:32
Request 38 is not allowed at 16:14:32
Request 39 is not allowed at 16:14:32
Request 40 is not allowed at 16:14:32
Request 41 is not allowed at 16:14:32
Request 42 is allowed at 16:14:32
Request 43 is not allowed at 16:14:33
Request 44 is not allowed at 16:14:33
Request 45 is not allowed at 16:14:33
Request 46 is allowed at 16:14:33
Request 47 is not allowed at 16:14:33
Request 48 is not allowed at 16:14:33
Request 49 is not allowed at 16:14:33
Request 50 is not allowed at 16:14:33
Request 51 is allowed at 16:14:33
Request 52 is not allowed at 16:14:34
Request 53 is not allowed at 16:14:34
Request 54 is not allowed at 16:14:34
Request 55 is not allowed at 16:14:34
Request 56 is allowed at 16:14:34
Request 57 is not allowed at 16:14:34
Request 58 is not allowed at 16:14:34
Request 59 is not allowed at 16:14:34
Request 60 is allowed at 16:14:34
Request 61 is not allowed at 16:14:35
Request 62 is not allowed at 16:14:35
从输出结果可以观察到每秒限制2个请求通过。
实现原理
假设这样实现令牌桶:
用阻塞队列模拟令牌桶,开一个定时器,定时队列里放令牌,使用生产者-消费者模式实现即可。
这个方式看起来好像没什么问题,但开启定时器需要新开线程,限流本就是在高并发场景下使用,额外开启线程会给系统带来更多开销。另外,假设我们是针对热点商户进行限流,如果有1万个热点商户,我们就要开启1万个定时器,这个开销是很大的。RateLimiter使用一种巧妙的方式,基于请求的间隔时间,来“模拟”出定时的效果。
基本思路是在每次请求处理完毕后,记录一个时间戳。然后,在下一次请求到达时,检查当前时间与上一次记录的时间戳之差,根据这个差值来决定是否执行定时任务或如何调整下一次的执行逻辑。
Limiter结构体:
type Limit float64
type Limiter struct {
mu sync.Mutex
limit Limit
burst int
tokens float64
// last is the last time the limiter's tokens field was updated
last time.Time
// lastEvent is the latest time of a rate-limited event (past or future)
lastEvent time.Time
}
对结构体字段的解释:
mu sync.Mutex
:这是一个互斥锁,用于在并发环境中保护Limiter
的状态,确保在多个goroutine同时访问或修改Limiter
的字段时不会出现数据竞争。limit Limit
:表示令牌桶算法中的速率限制。这个值决定了每秒钟可以向桶中添加多少个令牌。例如:limit
是 1.0,那么每秒可以向桶中添加 1 个令牌。burst int
:令牌桶的最大容量,即桶中最多可以存储多少个令牌,系统可以允许的最大突发请求数。tokens float64
:当前令牌桶中的令牌数量。last time.Time
:上一次更新tokens
字段的时间戳。lastEvent time.Time
:最近一次受速率限制的事件的时间戳。
核心函数advance
函数被命名为 advance
是因为有“前进”或“推进”的意思,在这个时间推进过程中应该增加的令牌数量的操作。
func (lim *Limiter) advance(t time.Time) (newT time.Time, newTokens float64) {
last := lim.last
// 处理时间回拨
if t.Before(last) {
last = t
}
// 计算从 last 到 t 经过的时间
elapsed := t.Sub(last)
// 计算在这段时间内根据令牌率(lim.limit)可以累积的新增令牌数
delta := lim.limit.tokensFromDuration(elapsed)
// 将当前令牌数(lim.tokens)与新增的令牌数(delta)相加,得到新的令牌数
tokens := lim.tokens + delta
// 检查并限制令牌数不超过令牌桶的容量
if burst := float64(lim.burst); tokens > burst {
tokens = burst
}
return t, tokens
}
reserveN函数
func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration) Reservation {
lim.mu.Lock()
defer lim.mu.Unlock()
if lim.limit == Inf {
return Reservation{
ok: true,
lim: lim,
tokens: n,
timeToAct: t,
}
} else if lim.limit == 0 {
var ok bool
if lim.burst >= n {
ok = true
lim.burst -= n
}
return Reservation{
ok: ok,
lim: lim,
tokens: lim.burst,
timeToAct: t,
}
}
t, tokens := lim.advance(t)
// Calculate the remaining number of tokens resulting from the request.
// 从当前时间窗口的令牌数中减去请求的令牌数n
tokens -= float64(n)
// Calculate the wait duration
// 如果当前可用的令牌不足以满足请求,调用者需要等待多长时间才能再次尝试或得到足够的令牌来执行操作。
var waitDuration time.Duration
if tokens < 0 {
waitDuration = lim.limit.durationFromTokens(-tokens)
}
// Decide result
// maxFutureReserve 限制了如果当前没有足够的令牌来满足请求,调用者愿意等待的最长时间。这有助于防止请求无限期地等待令牌
// maxFutureReserve设置为0,ok为false(tokens < 0的情况下)
ok := n <= lim.burst && waitDuration <= maxFutureReserve
// Prepare reservation
r := Reservation{
ok: ok,
lim: lim,
limit: lim.limit,
}
if ok {
r.tokens = n
r.timeToAct = t.Add(waitDuration)
// Update state
lim.last = t
lim.tokens = tokens
lim.lastEvent = r.timeToAct
}
return r
}
当maxFutureReserve=0,tokens >= 0:
ok := n <= lim.burst && 0 <= 0
ok为true
当ok为true,更新limit状态。
if ok {
r.tokens = n
r.timeToAct = t.Add(waitDuration)
// Update state
lim.last = t
lim.tokens = tokens
lim.lastEvent = r.timeToAct
}
ReserveN函数
ReserveN返回一个Reservation,指示调用者必须等待多长时间才能发生n个事件。
如果n超过Limiter的突发大小,则返回的Reservation的OK()方法返回false。
func (lim *Limiter) ReserveN(t time.Time, n int) *Reservation {
r := lim.reserveN(t, n, InfDuration)
return &r
}
这里maxFutureReserve = InfDuration
reserveN函数ok的判断逻辑:
ok := n <= lim.burst && waitDuration <= maxFutureReserve
waitDuration <= maxFutureReserve永远为true。
// 创建一个每秒允许2个令牌的Limiter,桶的容量为10
limiter := rate.NewLimiter(rate.Every(500*time.Millisecond), 10)
for i := 0; i < 200; i++ {
r := limiter.ReserveN(time.Now(), 1)
if !r.OK() {
// Not allowed to act! Did you remember to set lim.burst to be > 0 ?
return
}
if i == 100 {
time.Sleep(10 * time.Second)
}
time.Sleep(r.Delay())
fmt.Println("Request", i, "is allowed at", time.Now().Format("15:04:05"))
}
Wait函数
这个方法是对另一个方法WaitN
的“快捷方式”,专门用于等待一个请求被允许通过限流器。
// Wait is shorthand for WaitN(ctx, 1).
func (lim *Limiter) Wait(ctx context.Context) (err error) {
return lim.WaitN(ctx, 1)
}
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
// The test code calls lim.wait with a fake timer generator.
// This is the real timer generator.
newTimer := func(d time.Duration) (<-chan time.Time, func() bool, func()) {
timer := time.NewTimer(d)
return timer.C, timer.Stop, func() {}
}
return lim.wait(ctx, n, time.Now(), newTimer)
}
WaitN阻塞直到lim允许n个事件发生。如果n超过限制器的突发大小、context被取消或预期等待时间超过context的截止日期,则返回错误。如果速率限制为Inf,则忽略突发限制。
// wait is the internal implementation of WaitN.
func (lim *Limiter) wait(ctx context.Context, n int, t time.Time, newTimer func(d time.Duration) (<-chan time.Time, func() bool, func())) error {
lim.mu.Lock()
burst := lim.burst
limit := lim.limit
lim.mu.Unlock()
if n > burst && limit != Inf {
return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst)
}
// Check if ctx is already cancelled
select {
case <-ctx.Done():
return ctx.Err()
default:
}
// Determine wait limit
waitLimit := InfDuration
if deadline, ok := ctx.Deadline(); ok {
waitLimit = deadline.Sub(t)
}
// Reserve
r := lim.reserveN(t, n, waitLimit)
if !r.ok {
return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
}
// Wait if necessary
delay := r.DelayFrom(t)
if delay == 0 {
return nil
}
ch, stop, advance := newTimer(delay)
defer stop()
advance() // only has an effect when testing
select {
case <-ch:
// We can proceed.
return nil
case <-ctx.Done():
// Context was canceled before we could proceed. Cancel the
// reservation, which may permit other events to proceed sooner.
r.Cancel()
return ctx.Err()
}
}
总结
- 本项目令牌桶的实现并没有单独起一个goroutine来添加令牌,是基于时间来实现的,这样令牌桶的更新是懒惰的(即只在需要时更新),不需要额外的goroutine或线程来定期检查并添加令牌,从而减少了系统的开销。