13. Go中常见限流算法示例代码

代码地址：
计数器法：https://gitee.com/lymgoforIT/golang-trick/tree/master/08-count-limit-rate
令牌桶算法：https://gitee.com/lymgoforIT/golang-trick/tree/master/09-token-bucket-limiter
漏桶算法：https://gitee.com/lymgoforIT/golang-trick/tree/master/10-leaky-bucket-limiter

计数器、滑动窗口、漏斗算法、令牌桶算法是我们常见的几个限流算法，虽然已经有现成的包可以直接使用，但是理解各算法的原理以及代码实现还是很有帮助的。

1. 计数器法

计数器是一种最简单限流算法，其原理就是：在一段时间间隔内，对请求进行计数，与阀值进行比较判断是否需要限流，一旦到了时间临界点，将计数器清零。

• 可以在程序中设置一个变量 count，当过来一个请求我就将这个数 +1，同时记录请求时间。
• 当下一个请求来的时候判断count的计数值是否超过设定的频次，以及当前请求的时间和第一次请求时间是否在 1 分钟内。
• 如果在 1 分钟内并且超过设定的频次则证明请求过多，后面的请求就拒绝掉。
• 如果该请求与第一个请求的间隔时间大于计数周期，且 count 值还在限流范围内，就重置 count。

常见使用redis。比如对每秒限流，可以每秒设置一个key（limit-2023-09-21-20:11:11、limit-2023-09-21-20:11:12）。首先获取当前时间拼接一个如上的key(即拼上自定义的前缀limit-)。如果key存在且未超过某个阈值就自增，超过阈值就拒绝；如果key不存在就代表这新的一秒没有请求则重置计数。

Go 实现

type CountLimiter struct {
	rate  int64         // 计数周期内运行的最大请求数
	begin time.Time     // 当前轮计数开始时间
	count int64         // 当前计数周期累计的请求数
	cycle time.Duration // 计数周期，如统计1秒内的总请求数，那么计数周期就是1秒
	lock  sync.Mutex    // 判断能否放行时需要加锁操作
}

func NewCountLimiter(rate int64, cycle time.Duration) *CountLimiter {
	return &CountLimiter{
		rate:  rate,
		begin: time.Now(),
		count: 0,
		cycle: cycle,
		lock:  sync.Mutex{},
	}
}

func (c *CountLimiter) Allow() bool {
	c.lock.Lock()
	defer c.lock.Unlock()
	if c.count == c.rate { // 达到最大限流时，判断时间是否也超过统计周期了，超了可以重置限流器了，没有超说明还在当前统计周期内，应该拦截
		if time.Now().Sub(c.begin) > c.cycle {
			c.Reset()
			return true
		} else {
			return false
		}
	} else { // 还没有达到最大限流数，那么不管时间周期是否超出统计计数周期，都可以放行
		c.count++
		return true
	}
}

func (c *CountLimiter) Reset() {
	c.begin = time.Now()
	c.count = 0
}

func main() {
	countLimiter := NewCountLimiter(3, time.Second) // 1s内不可超过3个请求
	var wg sync.WaitGroup
	wg.Add(10)
	for i := 0; i < 10; i++ {
		go func(i int) {
			defer wg.Done()
			if countLimiter.Allow() {
				fmt.Println(fmt.Sprintf("当前时间%s,请求%d通过了", time.Now().String(), i))
			} else {
				fmt.Println(fmt.Sprintf("当前时间%s,请求%d被限流了", time.Now().String(), i))
			}
		}(i)
		time.Sleep(200 * time.Millisecond)
	}
	wg.Wait()
}

缺点

如果有个需求对于某个接口 /query 每分钟最多允许访问 200 次，假设有个用户在第 59 秒的最后几毫秒瞬间发送 200 个请求，当 59 秒结束后 Counter 清零了，他在下一秒的时候又发送 200 个请求。那么在 1 秒钟内这个用户发送了 2 倍的请求，这个是符合我们的设计逻辑的，这也是计数器方法的设计缺陷，系统可能会承受恶意用户的大量请求，甚至击穿系统。

2. 令牌桶算法

算法简介

令牌桶算法（Token Bucket）是网络流量整形（Traffic Shaping）和速率限制（Rate Limiting）中最常使用的一种算法。典型情况下，令牌桶算法用来控制发送到网络上的数据的数目，并允许突发数据的发送。想象有一个木桶，以固定的速度往木桶里加入令牌，木桶满了则不再加入令牌。服务收到请求时尝试从木桶中取出一个令牌，如果能够得到令牌则继续执行后续的业务逻辑；如果没有得到令牌，直接返回访问频率超限的错误码或页面等，不继续执行后续的业务逻辑

由于木桶内只要有令牌，请求就可以被处理，所以令牌桶算法可以支持突发流量。同时由于往木桶添加令牌的速度是固定的，且木桶的容量有上限，所以单位时间内处理的请求数目也能够得到控制，起到限流的目的。假设加入令牌的速度为 1token/10ms(则1s内最多放置100个令牌，因此QPS期望是100左右），另一方面，桶的容量为500，在请求比较的少的时候（小于每10毫秒1个请求）时，木桶可以先"攒"一些令牌（最多500个）。当有突发流量时，一下把木桶内的令牌取空，也就是有500个在并发执行的业务逻辑，之后要等每10ms补充一个新的令牌才能接收一个新的请求。

木桶的容量设置：需要考虑业务逻辑的资源消耗和机器能承载并发处理多少业务逻辑。
生成令牌的速度设置 ：太慢的话起不到“攒”令牌应对突发流量的效果，可根据预估或压测的QPS进行设置。

• 令牌按固定的速率被放入令牌桶中
• 桶中最多存放 B 个令牌，当桶满时，新添加的令牌被丢弃或拒绝
• 如果桶中的令牌不足 1个，则不会删除令牌，且请求将被限流（丢弃或阻塞等待）

令牌桶限制的是平均流入速率（允许突发请求，只要有令牌就可以处理，支持一次拿3个令牌，4个令牌…），并允许一定程度突发流量。

适用场景

适合电商抢购或者微博出现热点事件这种场景，因为在限流的同时可以应对一定的突发流量。如果采用均匀速度处理请求的算法，在发生热点时间的时候，会造成大量的用户无法访问，对用户体验的损害比较大。

Go 实现

假设设置每100ms生产一个令牌，记录最近一次访问的时间戳 lastTime 和令牌数，每次请求时如果 now - lastTime > 100ms, 增加 (now - lastTime) / 100ms个令牌。然后，如果令牌数 > 0，令牌数 -1 继续执行后续的业务逻辑，否则返回请求频率超限的错误码或页面。

上面的算法是对整体的请求进行的限流，如果是要对用户或IP进行限流，则可以使用map[string]Limiter控制，key为userId或IP，value为对应的限流器。

package main

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

type TokenBucketLimiter struct {
	lock     sync.Mutex
	rate     time.Duration // 多长时间放入一个令牌，即放入令牌的速率
	capacity int64         // 令牌桶的容量，控制最多放入多少令牌，也即突发最大并发量
	tokens   int64         // 当前桶中已有的令牌数量
	lastTime time.Time     // 上次放入令牌的时间，避免开启协程定时去放入令牌，而是请求到来时懒加载的方式(now - lastTime) / rate放入令牌
}

func NewTokenBucketLimiter(rate time.Duration, capacity int64) *TokenBucketLimiter {
	if capacity < 1 {
        panic("token bucket capacity must be large 1")
    }
	return &TokenBucketLimiter{
		lock:     sync.Mutex{},
		rate:     rate,
		capacity: capacity,
		tokens:   0,
		lastTime: time.Time{},
	}
}

func (tbl *TokenBucketLimiter) Allow() bool {
	tbl.lock.Lock() // 加锁避免并发错误
	defer tbl.lock.Unlock()

	// 如果 now 与上次请求的间隔超过了 token rate
	// 则增加令牌，更新lastTime
	now := time.Now()
	if now.Sub(tbl.lastTime) > tbl.rate {
		tbl.tokens += int64((now.Sub(tbl.lastTime)) / tbl.rate) // 放入令牌
		if tbl.tokens > tbl.capacity {
			tbl.tokens = tbl.capacity // 总令牌数不能大于桶的容量
		}
		tbl.lastTime = now // 更新上次往桶中放入令牌的时间
	}

	if tbl.tokens > 0 { // 令牌数是否充足
		tbl.tokens -= 1
		return true
	}

	return false // 令牌不足，拒绝请求
}

func main() {
	tbl := NewTokenBucketLimiter(10, 5) // 每10ms放一个令牌，1s放100个，桶容量（最大突发流量）为5
	for i := 0; i < 10; i++ {
		fmt.Println(tbl.Allow())  // 模拟突发流量10个请求，超过桶容量5
	}
	time.Sleep(100 * time.Millisecond)
	fmt.Println(tbl.Allow())
}

漏桶算法

算法简介

与令牌桶是“反向”的算法，当有请求到来时先放到木桶中，worker以固定的速度从木桶中取出请求进行响应。如果木桶已经满了，直接返回请求频率超限的错误码或者页面。

漏桶算法有以下特点：

• 漏桶具有固定容量，出水速率是固定常量（流出请求）
• 如果桶是空的，则不需流出水滴
• 可以以任意速率流入水滴到漏桶（流入请求）
• 如果流入水滴超出了桶的容量，则流入的水滴溢出（新请求被拒绝）
• 漏桶限制的是常量流出速率（即流出速率是一个固定常量值），所以最大的速率就是出水的速率，不能出现突发流量。

适用场景

流量最均匀的限流方式，一般用于流量“整形”，例如保护数据库的限流。先把对数据库的访问加入到木桶中，worker再以db能够承受的qps从木桶中取出请求，去访问数据库。不太适合电商抢购和微博出现热点事件等场景的限流，一是应对突发流量不是很灵活，二是如果需要对用户或者IP限流，则用map为每个user_id/ip维护一个队列(木桶)，worker从这些队列中拉取任务，资源的消耗会比较大。

Go 实现

通常使用队列来实现，在go语言中可以通过带缓冲的通道buffered channel来快速实现，任务加入channel，开启一定数量的 worker 从 channel 中获取任务执行,这一定数量的worker表示的就是限流，如开启5个worker，便表示限流速度为5,即同一时刻最多处理5个请求。

package main

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

type Result struct {
	Msg string // 根据实际情况定义返回结果需要哪些字段
}

type Handler func() Result // 处理函数的形式也应该根据具体需要而定

type Task struct {
	id      int64       // 任务id
	result  chan Result // 任务的执行结果，即请求的响应结果
	handler Handler     // 请求的执行函数
}

func NewTask(id int, handler Handler) Task {
	return Task{
		handler: handler,
		result:  make(chan Result),
		id:      int64(id),
	}
}

type LeakyBucketLimiter struct {
	bucketSize int64     // 桶的大小
	workerNum  int64     // 工作者数量，即最大并发数
	taskChan   chan Task // 用于存放请求
}

func NewLeakyBucketLimiter(bucketSize, workerNum int64) *LeakyBucketLimiter {
	if capacity < 1 {
        panic("capacity must be large 1")
    }
    if workerNum < 1 {
        panic("workerNum must be large 1")
    }
	return &LeakyBucketLimiter{
		bucketSize: bucketSize,
		workerNum:  workerNum,
		taskChan:   make(chan Task, bucketSize),
	}
}

func (lbl *LeakyBucketLimiter) AddTask(task Task) bool { // 类似其他限流算法的Allow方法
	// 如果木桶已经满了，或者任务执行失败或超时了，返回false
	select {
	case lbl.taskChan <- task:  // 利用了select的特性判断是否能往通道中添加任务
	default:
		fmt.Printf("请求%d被拒绝了\n", task.id)
		return false
	}

	// 如果成功入桶，调用者会等待Task的Handler执行结果
	// 由于Task的result是无缓冲的通道，不应该让其无限等待阻塞，否则出现问题时，不往该chan写，就会一直阻塞在这里了，泄漏
	// 因此设置一个超时时间
	//resp := <-task.result
	//fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
	select {
	case resp := <-task.result:
		fmt.Printf("请求%d运行成功，结果为：%v\n", task.id, resp)
	case <-time.After(5 * time.Second): // 超时时间可以稍微设置长一点点，因为任务放入桶中后，可能需要排队一点时间才被拉取出来执行
		return false // 这里超时当被限流处理
	}

	return true
}

func (lbl *LeakyBucketLimiter) Start(ctx context.Context) {
	// 开启workerNum个协程从木桶拉取任务执行
	for i := 0; int64(i) < lbl.workerNum; i++ {
		go func(ctx context.Context) {
			defer func() { // 铁则：开启的子协程一定要捕获异常，否则一旦出现异常不捕获，会导致程序退出
				if err := recover(); err != any(nil) {
					fmt.Println("捕获到异常")
				}
			}()

			for { // 持续监听，拉取任务执行
				select {
				case <-ctx.Done():
					fmt.Println("退出工作")
					return
				default:
					task := <-lbl.taskChan
					result := task.handler()
					task.result <- result // 处理结果写入对应Task的结果通道
				}
			}
		}(ctx)
	}
}

func main() {
	bucket := NewLeakyBucketLimiter(10, 4)
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	bucket.Start(ctx) // 开启消费者
	// 模拟20个并发请求
	var wg sync.WaitGroup
	wg.Add(20)
	for i := 0; i < 20; i++ {
		go func(id int) {
			defer wg.Done()
			task := NewTask(id, func() Result { // 这里的func应该根据实际需要定义为handler，写具体的业务逻辑和返回的result
				time.Sleep(300 * time.Millisecond)  // 模拟业务逻辑消耗的时间
				return Result{}
			})
			bucket.AddTask(task)  // 请求入桶
		}(i)
	}
	wg.Wait()
	time.Sleep(10 * time.Second)
}

上面的代码其实已经有点类似生成者消费者模型了，添加任务即往队列中放任务（生产消息），然后安排一定数量的协程（消费者）拉取任务执行（消费消息）