如何使用Golang通道实现一个简单的并发限流器

package main

import (
    "fmt"
    "sync"
    "time"
)

// RateLimiter 是一个基于通道的简单限流器
type RateLimiter struct {
    tokens   chan struct{} // 令牌通道
    capacity int           // 令牌桶容量
    fillRate time.Duration // 补充令牌的间隔时间 (例如，每 500ms 补充一个令牌)
    stop     chan struct{} // 用于停止令牌补充goroutine的信号
    once     sync.Once     // 确保 Stop() 只被调用一次
}

// NewRateLimiter 创建一个新的限流器实例
// capacity: 令牌桶的容量，决定了可以处理的瞬时最大请求数。
// fillRate: 补充一个令牌所需的时间。例如，time.Millisecond * 500 表示每 500ms 补充一个令牌，即每秒 2 个令牌。
func NewRateLimiter(capacity int, fillRate time.Duration) *RateLimiter {
    if capacity <= 0 {
        capacity = 1 // 确保容量至少为1
    }
    if fillRate <= 0 {
        fillRate = time.Millisecond * 100 // 设置一个默认的合理填充速率
    }

    limiter := &RateLimiter{
        tokens:   make(chan struct{}, capacity),
        capacity: capacity,
        fillRate: fillRate,
        stop:     make(chan struct{}),
    }

    // 初始填充令牌，使限流器在启动时即可处理请求
    for i := 0; i < capacity; i++ {
        limiter.tokens <- struct{}{}
    }

    go limiter.fillTokens() // 启动令牌补充goroutine
    return limiter
}

// fillTokens 持续向令牌桶中添加令牌
func (rl *RateLimiter) fillTokens() {
    ticker := time.NewTicker(rl.fillRate)
    defer ticker.Stop() // 确保在函数退出时停止定时器

    for {
        select {
        case <-ticker.C:
            // 定时器触发，尝试添加一个令牌
            select {
            case rl.tokens <- struct{}{}:
                // 令牌添加成功
            default:
                // 通道已满，无法添加更多令牌。这表示令牌桶已达到容量上限，无需额外操作。
            }
        case <-rl.stop:
            // 收到停止信号，退出goroutine
            fmt.Println("Rate limiter stopping token replenishment.")
            return
        }
    }
}

// Allow 阻塞式地获取一个令牌。如果令牌桶为空，则会一直阻塞直到有新令牌可用。
func (rl *RateLimiter) Allow() {
    <-rl.tokens // 从通道中接收一个令牌，如果通道为空则阻塞
}

// TryAllow 非阻塞式地尝试获取一个令牌。
// 如果成功获取到令牌，返回 true；否则（令牌桶为空），立即返回 false。
func (rl *RateLimiter) TryAllow() bool {
    select {
    case <-rl.tokens:
        return true
    default:
        return false
    }
}

// Stop 优雅地关闭限流器，停止后台的令牌补充goroutine。
func (rl *RateLimiter) Stop() {
    rl.once.Do(func() {
        close(rl.stop) // 关闭 stop 通道，通知 fillTokens goroutine 退出
    })
}

func main() {
    // 示例用法：创建一个每秒允许 2 个请求的限流器
    // 容量为 2，意味着可以处理 2 个瞬时请求。
    // 填充速率为 500ms，表示每 500ms 补充一个令牌，即每秒 2 个令牌。
    limiter := NewRateLimiter(2, time.Millisecond*500)
    defer limiter.Stop() // 确保在程序退出时关闭限流器

    var wg sync.WaitGroup
    start := time.Now()

    fmt.Println("Starting requests...")

    // 模拟 10 个请求
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            limiter.Allow() // 请求会在这里等待，直到获取到令牌
            fmt.Printf("Request %d processed at %v\n", id, time.Since(start).Round(time.Millisecond))
        }(i)
        // 稍微延迟一下，模拟请求并非同时到达，否则所有请求可能会同时阻塞
        time.Sleep(time.Millisecond * 100)
    }

    wg.Wait() // 等待所有请求完成
    fmt.Printf("All requests finished in %v\n", time.Since(start).Round(time.Millisecond))
}

Allow()

TryAllow()

sync.Mutex

sync.Cond

<-limiter.tokens

<-ch

select { case <-ch: ... default: ... }

fillTokens

capacity

fillRate

Stop()

Stop()

Allow()

Stop

tokens

context.Context