[go]调用接口时进行限流

最近有一个需求：在我写的后端调用算法接口，但是要限流调用，不能接收到请求，处理好后，直接调用算法接口。

也就是要进行限流处理

最后询问了同事，解决方案是做一个令牌桶，在项目启动时就开始监听处理好的任务，但是在一个指定的时间周期内，只能发送指定次数的算法请求。

实现方案如下：

// TokenBucket 代表一个令牌桶，用于控制流量。
type TokenBucket struct {
    ch        chan struct{}
    tokens    int
    maxTokens int
    interval  time.Duration
    closeChan chan struct{}
    wg        sync.WaitGroup
}

// NewTokenBucket 创建一个新的令牌桶，指定令牌的最大数量和令牌添加的时间间隔。
func NewTokenBucket(maxTokens int, interval time.Duration) *TokenBucket {
    return &TokenBucket{
       ch:        make(chan struct{}),
       tokens:    0,
       maxTokens: maxTokens,
       interval:  interval,
       closeChan: make(chan struct{}),
    }
}

// Start 开始令牌桶的工作，它会定期向桶中添加令牌。
func (tb *TokenBucket) Start() {
    zap.L().Info("令牌桶开始工作")
    tb.wg.Add(1)
    go func() {
       defer tb.wg.Done()
       ticker := time.NewTicker(tb.interval)
       defer ticker.Stop()

       for {
          select {
          case <-ticker.C:
             tb.addToken()
          case <-tb.closeChan:
             return
          }
       }
    }()
}

// addToken 向令牌桶中添加一个令牌。
func (tb *TokenBucket) addToken() {
    if tb.tokens >= tb.maxTokens {
       tb.tokens = tb.maxTokens // 限制令牌数不超过最大值
       zap.L().Info("令牌桶已满，无法添加更多令牌")
    } else {
       tb.tokens++
       tb.ch <- struct{}{}
    }
}

// Pour 尝试从令牌桶中取出一个令牌来执行任务，如果没有令牌可用，则等待。
func (tb *TokenBucket) Pour() {
    <-tb.ch // 等待有令牌可用
    tb.tokens--
}

// Close 停止令牌桶的工作。
func (tb *TokenBucket) Close() {
    close(tb.closeChan)
    tb.wg.Wait()
}

👉👈大佬们还有没有其他好的实现方法，可以发在评论区。