Go 实现 RateLimiter 的代码示范

11 min read

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() 等待它们完成后退出程序。