实现一个单机版令牌桶限流器（字节）

令牌桶说明

令牌桶是一种常见的限流算法，它的工作原理如下：

令牌生成：系统以固定的速率向桶中添加令牌
令牌消耗：请求需要消耗令牌才能被处理
容量限制：桶有最大容量，多余的令牌会被丢弃
请求处理：当有请求到来时，如果有足够的令牌则处理，否则拒绝

参考golang实现

go 复制代码

package main

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

// TokenBucket 令牌桶结构
type TokenBucket struct {
	capacity  int64        // 桶容量
	tokens    int64        // 当前令牌数
	rate      int64        // 每秒生成令牌速率
	mutex     sync.Mutex   // 互斥锁
	ticker    *time.Ticker // 定时器
	stopChan  chan bool    // 停止信号
}

// NewTokenBucket 创建一个新的令牌桶
func NewTokenBucket(capacity, rate int64) *TokenBucket {
	tb := &TokenBucket{
		capacity: capacity,
		tokens:   capacity, // 初始时桶是满的
		rate:     rate,
		ticker:   time.NewTicker(time.Second), // 每秒触发一次
		stopChan: make(chan bool),
	}
	
	// 启动后台goroutine定时生成令牌
	go func() {
		for {
			select {
			case <-tb.ticker.C:
				tb.mutex.Lock()
				// 添加令牌，不超过桶容量
				tb.tokens = min(tb.tokens+tb.rate, tb.capacity)
				tb.mutex.Unlock()
			case <-tb.stopChan:
				tb.ticker.Stop()
				return
			}
		}
	}()
	
	return tb
}

// Allow 判断是否允许执行操作
func (tb *TokenBucket) Allow() bool {
	return tb.AllowN(1)
}

// AllowN 判断是否允许执行n个操作
func (tb *TokenBucket) AllowN(n int64) bool {
   // 只负责检查和消费令牌，减少了锁的持有时间，提高了并发性能
	tb.mutex.Lock()
	defer tb.mutex.Unlock()
	
	// 判断令牌是否足够
	if tb.tokens >= n {
		tb.tokens -= n
		return true
	}
	
	return false
}

// Stop 停止令牌桶
func (tb *TokenBucket) Stop() {
	close(tb.stopChan)
}

// min 返回较小值
func min(a, b int64) int64 {
	if a < b {
		return a
	}
	return b
}

func main() {
	// 创建一个容量为10，每秒生成2个令牌的令牌桶
	bucket := NewTokenBucket(10, 2)
	defer bucket.Stop() // 程序结束前停止令牌桶，优雅关闭，避免协程泄漏
	
	fmt.Println("开始测试令牌桶限流器...")
	
	// 测试1: 连续快速请求
	fmt.Println("\n测试1: 连续快速请求")
	for i := 0; i < 15; i++ {
		if bucket.Allow() {
			fmt.Printf("请求%d: 通过\n", i+1)
		} else {
			fmt.Printf("请求%d: 被拒绝\n", i+1)
		}
		time.Sleep(100 * time.Millisecond) // 间隔100毫秒发送请求，15个请求在1.5s内发出，共需要15个令牌
	}
	
	// 等待一段时间让令牌桶重新填满
	fmt.Println("\n等待5秒让令牌桶重新填充...")
	time.Sleep(5 * time.Second)
	
	// 测试2: 短时间大量请求
	fmt.Println("\n测试2: 短时间大量请求")
	passed := 0
	for i := 0; i < 15; i++ {
		if bucket.Allow() {
			passed++
			fmt.Printf("请求%d: 通过\n", i+1)
		} else {
			fmt.Printf("请求%d: 被拒绝\n", i+1)
		}
	}
	fmt.Printf("总共通过: %d/15\n", passed)
	
	// 测试3: 批量请求
	fmt.Println("\n测试3: 批量请求测试")
        // 创建一个容量为20，每秒生成5个令牌的令牌桶
	bucket2 := NewTokenBucket(20, 5)
        // 优雅关闭
	defer bucket2.Stop()
	
	// 请求5个令牌
	if bucket2.AllowN(5) {
		fmt.Println("批量请求5个令牌: 通过")
	} else {
		fmt.Println("批量请求5个令牌: 被拒绝")
	}
	
	// 请求15个令牌（超过当前令牌数）
	if bucket2.AllowN(15) {
		fmt.Println("批量请求15个令牌: 通过")
	} else {
		fmt.Println("批量请求15个令牌: 被拒绝")
	}
	
	// 等待一段时间后再尝试
	time.Sleep(3 * time.Second)
	if bucket2.AllowN(15) {
		fmt.Println("等待3s后批量请求15个令牌: 通过")
	} else {
		fmt.Println("等待3s后批量请求15个令牌: 被拒绝")
	}
}