Golang | 令牌桶限流算法

限流算法的目的是控制对后端接口的访问频率，防止因过度访问导致系统崩溃。
网站登录接口是限流的典型例子，爬虫或恶意用户可能疯狂调用登录接口，导致数据库压力过大。
通过限制接口的QPS（每秒查询率），可以保护后端数据库不受冲击。
令牌桶算法是一种经典的限流算法，将请求想象成令牌，桶以一定速度放入令牌。
桶的容量决定瞬间最多能满足的请求数，生产速度代表平均供应速度。
谷歌官方提供了基于令牌桶算法的实现，位于guava库的x包下。
使用rate.newLimit方法创建限流器，指定时间间隔和桶容量。
限流器接受三种取令牌方式：wait、allow和reserve。
wait方式会等待桶内令牌足够为止；allow方式立即返回是否允许取走令牌；reserve方式会计算等待时间。

go 复制代码

package tests

import (
	"sync/atomic"
	"time"
	"golang.org/x/time/rate"
)

// TotalQuery 记录请求总数
var TotalQuery int32

// Handler 模拟接口
func Handler() {
	atomic.AddInt32(&TotalQuery, 1)
	time.Sleep(50 * time.Millisecond)
}

func CallHandler() {
	// 每隔100ms生成一个令牌，最大QPS限制为10
	limiter := rate.NewLimiter(rate.Every(100*time.Millisecond), 10)
	for {
		// 方式一：WaitN()
		//ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
		//defer cancel()
		//limiter.WaitN(ctx, 1) // 阻塞，直到桶中有N个令牌。N=1时等价于Wait(ctx)
		//Handler()

		// 方式二：AllowN()
		// 不等，当前桶中是否至少还有N个令牌，如果有则返回true。N=1时等价于Allow(time.Time)
		//if limiter.AllowN(time.Now(), 1) {
		//	Handler()
		//}

		// 方式三：ReserveN()
		// reserve.Delay()告诉你还需要等多久才会有充足的令牌，等待相应时间，再执行
		reserve := limiter.ReserveN(time.Now(), 1)
		time.Sleep(reserve.Delay())
		Handler()
	}
}

令牌桶（Token Bucket）算法是一种常用的限流算法，常用于网络流量控制和高并发服务的请求限速。它本质上是一个允许突发请求但控制平均速率的机制，兼具了平滑输出和灵活突发的优点。
令牌桶算法通过一个"桶"来存储"令牌"，系统按照固定速率向桶中添加令牌，处理请求时必须"取走"一个令牌，如果桶里没有令牌了，请求就被拒绝或延迟处理。
每次处理请求时：
- 令牌补充：根据当前时间与上次补充时间的差值，计算应添加的令牌数量：tokens += (now - last_refill_time) * r
- 令牌溢出处理：如果 tokens > b，令牌数取为桶最大容量。
- 判断请求是否能通过：
  - 如果 tokens >= 1，扣除一个令牌，允许请求通过。
  - 否则，拒绝请求或阻塞等待。