【Go】令牌桶限流算法

package main
 
import (
    "fmt"
    "time"
)
 
// TokenBucketLimiter 表示令牌桶限流器
type TokenBucketLimiter struct {
    rate         int           // 令牌产生的速率（每秒产生令牌的数量）
    tokens       int           // 当前持有的令牌数
    lastTime     time.Time     // 上次更新令牌的时间
    maxTokens    int           // 最大令牌数
    tokenChannel chan struct{} // 用于同步的通道
}
 
// NewTokenBucketLimiter 创建一个新的令牌桶限流器
func NewTokenBucketLimiter(rate int, maxTokens int) *TokenBucketLimiter {
    return &TokenBucketLimiter{
        rate:         rate,
        tokens:       maxTokens,
        lastTime:     time.Now(),
        maxTokens:    maxTokens,
        tokenChannel: make(chan struct{}, maxTokens),
    }
}
 
// Wait 等待获取令牌
func (l *TokenBucketLimiter) Wait() {
    // 添加令牌
    l.addTokens()
 
    // 尝试获取令牌
    select {
    case l.tokenChannel <- struct{}{}:
        // 成功获取令牌，继续执行
    default:
        // 无法获取令牌，等待或抛出错误
        time.Sleep(100 * time.Millisecond)
        l.Wait() // 递归等待
    }
}
 
// addTokens 添加新的令牌到令牌桶中
func (l *TokenBucketLimiter) addTokens() {
    now := time.Now()
    elapsed := now.Sub(l.lastTime).Seconds()
    l.lastTime = now
    toAdd := int(elapsed * l.rate)
 
    if toAdd > 0 {
        l.tokens += toAdd
        if l.tokens > l.maxTokens {
            l.tokens = l.maxTokens
        }
    }
}
 
func main() {
    // 创建限流器，速率为每秒1个令牌
    limiter := NewTokenBucketLimiter(1, 2)
 
    // 模拟5个并发请求
    for i := 0; i < 5; i++ {
        go func() {
            for {
                limiter.Wait() // 等待获取令牌
                fmt.Println("处理请求")
                time.Sleep(500 * time.Millisecond) // 模拟请求处理时间
            }
        }()
    }
 
    // 主线程无限循环，模拟服务运行
    for {
        time.Sleep(1 * time.Second)
    }
}

这段代码实现了一个简单的令牌桶限流器，并在main函数中通过模拟并发请求的方式展示了其使用方式。它定义了TokenBucketLimiter结构体，并提供了创建新限流器和等待获取令牌的方法。addTokens函数负责根据时间添加令牌到桶中。在main函数中，我们创建了一个限流器，并启动了5个并行的goroutine来模拟并发请求，每个请求都会在执行前等待获取令牌。