掌握流量管理：深入Ratelimiting Queue的结构与源码解析

站长

2024年03月22日 16:44 · 阅读数 4

我是 LEE，老李，一个在 IT 行业摸爬滚打 17 年的技术老兵。

在开发经验中 RateLimiting Queue 的使用频率并不高，甚至比 Delaying Queue 和 Priority Queue 低了许多，但是有一些场景非常的需要，比如在 Kubernetes 开发一个 Controller 需要处理大量的数据，但是又不想一次性处理太多，这个时候就可以使用 RateLimiting Queue 来限制每次处理的数据量。还有就是比如一些大型的 API 服务，为了保证服务的稳定性，也会使用 RateLimiting Queue 来限制每秒的请求量。

RateLimiting Queue 整体设计采用了跟 client-go 项目中 RateLimiting Queue 一样的设计，但是在具体的实现上有所差异。RateLimiting Queue 的实现是基于 Delaying Queue 的，Delaying Queue 的实现是基于 Queue 的。通俗的说：RateLimiting Queue 是基于 Queue + Min Quad Heap + Limiter 组合而成。

WorkQueue 项目开源地址在这里：github.com/shengyanli1…

如果你还没有看过 WorkQueue 的前几篇文章，建议先看一下，方便你对后续内容的理解：

事件背景

年底了，开始一个核心项目的修复和完善。这个事情我牵头，然后还有一个同事协助我完成。在修复和完善的过程中，我们发现了一个非常严重的问题，原有代码的中的延迟队列部分是一个 channel 存放着元素，然后通过一个 goroutine 来处理这个 channel 中的元素，而这个 goroutine 的 handleFunc 里面的处理逻辑及其复杂，并通过一系列手段来控制元素处理的速度。同时原有代码使用了 golang 原生库 golang.org/x/time/rate, 混合使用了 Wait 和 Allow，看着头就晕。

实际想想，要实现一个 RateLimiting Queue 实际非常简单，为什么要用这么复杂控制逻辑，而且还有可能会出现 panic。于是我就想着把这个限速代码用 WorkQueue 项目代还，这样就可以避免这个问题了。

当然翻别人的代码是存在风险，但是我还是决定这么做，因为我相信 WorkQueue 项目的代码质量。

在这篇文章不会介绍翻代码的过程，只会介绍 WorkQueue 项目中 RateLimiting Queue 的实现原理和使用方法。

为什么要依赖 Delaying Queue

在介绍 RateLimiting Queue 一定避不开的问题，为什么要依赖 Delaying Queue。因为 Delaying Queue 有一个非常重要的特性，就是可以按照时间顺序来处理数据。 RateLimiting Queue 的实现就是基于这个特性来实现的。

我们尝试的思考下，当一组(量不小)的无序请求或者工作任务，要需要按照一定速度处理，我们需要怎么做才能保证这个速度呢？ 一个最简单的思路：我们可以把这些请求或者工作任务按照时间顺序排列，然后按照一定的速度来处理。

这个思路是不是很简单? 但是实现起来就有点复杂了。因为我们需要维护一个有序的队列，然后按照一定的速度来处理。这个时候 Delaying Queue 就派上用场了，它可以帮我们维护一个有序的队列，然后按照一定的速度来处理。

我想之前要被修复和完善的项目，应该是没有使用这个思路，所以才会写出那么复杂的代码。

（★）大叔白话: 这个就像现在我们过年了，要去车站坐车。我们随机的陆陆续续到车站，需要提前买票，然后在车站等待，等到了时间就可以上车了。这个过程就是 Delaying Queue 的实现原理。

Queue 介绍

RateLimiting Queue 凸显的就是一个 RateLimiting 的作用，那么我们怎么将 RateLimiting 和 Delaying Queue 结合起来呢？

参考了现在很多开源实现，觉得 Kubernetes 中 client-go 项目里面的 RateLimiting Queue 实现比较好，所以就参考了这个实现。利用 golang 原生库 golang.org/x/time/rate 中的 Reserve() 方法返回一个对象，再调用 Delay() 方法获得一个 time.Duration 的延迟时间，然后将这个延迟时间作为 Delaying Queue 的延迟时间，这样就可以实现 RateLimiting Queue 了。

提示: 调用该对象的 Delay() 方法就可以返回了需要等待的时间。如果等待时间为 0，则说明不用等待。

RateLimiting Queue

RateLimiting Queue 简介

RateLimiting Queue 是一个限速队列，它的特点是：投递的元素会按时间顺序，并根据设定的速率来执行。由于 RateLimiting Queue 的实现是基于 Delaying Queue 的，Delaying Queue 的实现是基于 Queue 的。所以说 RateLimiting Queue 包含了 Queue 和 Delaying Queue 的所有特性，同时还有 RateLimiting 的特性。

在 RateLimiting Queue 不需要关注元素的延迟时间或者权重，因为 RateLimiting Queue 内部有一个限速器，他会根据设定的速率来计算出延迟时间。并将元素运行延迟时间作为 Delaying Queue 的延迟时间，添加到 Delaying Queue 中。

在 RateLimiting Queue 如何实现一个高效的 Limiter？实际 golang.org/x/time/rate 已经帮我们实现了一个不错的 Limiter，是基于 TokenBucket 算法。如果你想基于 LeakyBucket 或者指数递增算法来实现一个 Limiter，就需要实现一些接口，让 Delaying Queue 能够调用。这个在后面的文章中会介绍。

新增接口:

AddLimited: 添加一个元素，需要对该元素进行限速处理。
Forget : 忘记一个元素，不需要对该元素进行限速处理。
NumLimitTimes : 返回一个元素被限速的次数。

// RateLimitingInterface 是 Queue 方法的接口
// RateLimitingInterface is the interface for Queue methods
type RateLimitingInterface interface {
	// 继承 DelayingQueue 接口
	// Inherit DelayingQueue
	DelayingInterface

	// AddLimited 添加一个元素，需要对该元素进行限速处理
	// AddLimited adds an element that needs to be rate-limited
	AddLimited(element any) error

	// Forget 忘记一个元素，不需要对该元素进行限速处理
	// Forget forgets an element that doesn't need to be rate-limited
	Forget(element any)

	// NumLimitTimes 返回一个元素被限速的次数
	// NumLimitTimes returns the number of times an element has been rate-limited
	NumLimitTimes(element any) int
}

新增 Callback 方法:

OnAddLimited : 添加元素后的回调
OnForget : 忘记元素后的回调
OnGetTimes : 返回一个元素被限速的次数

// RateLimitingCallback 是 Queue 的回调接口
// RateLimitingCallback is the callback interface for Queue
type RateLimitingCallback interface {
	// 继承 DelayingCallback 接口
	// Inherit DelayingCallback
	DelayingCallback

	// OnAddLimited 添加元素后的回调
	// OnAddLimited is the callback after adding an element
	OnAddLimited(any)

	// OnForget 忘记元素后的回调
	// OnForget is the callback after forgetting an element
	OnForget(any)

	// OnGetTimes 返回一个元素被限速的次数
	// OnGetTimes returns the number of times an element has been rate-limited
	OnGetTimes(any, int)
}

RateLimiting Queue 实现原理

RateLimiting Queue 整体结构设计如下：

掌握流量管理：深入Ratelimiting Queue的结构与源码解析

RateLimiting Queue 使用举例

当然 RateLimiting Queue 使用起来也非常简单，没有太多复杂的参数初始化过程。前提是你得限速器 Limiter 实现得足够的好，就可以获得你想要的效果。

还有就是和 Queue 是一样的，就是不论使用 Get 还是 GetWithBlock 之后一定要记得使用 Done 方法，来标记这个数据已经处理完成，否则再往 RateLimiting Queue 添加相同的数据的时候，会返回 ErrorQueueElementExist 错误。如果使用 Simple Queue，则不需要使用 Done 方法。

package main

import (
	"fmt"
	"time"

	"github.com/shengyanli1982/workqueue"
)

func main() {
	conf := NewRateLimitingQConfig().WithLimiter(NewBucketRateLimiter(float64(1), 1)).WithCallback(&rateLimitingcallback{})
	q := NewRateLimitingQueue(conf) // create a queue
	defer q.Stop()

	go func() {
		for {
			element, err := q.Get() // get element from queue
			if err != nil {
				fmt.Println(err)
				return
			}
			fmt.Println(">> get element:", element)
			q.Done(element) // mark element as done, 'Done' is required after 'Get'
		}
	}()

	_ = q.Add("hello") // add element to queue, immediately execute
	_ = q.Add("world")

	_ = q.AddAfter("delay", time.Second * 1) // add element to queue, execute after 1 second

	_ = q.AddLimited("burst") // add element with limit to queue, execute with rate limit(1/s, burst 1)
	_ = q.AddLimited("limit") // add element with limit to queue, execute with rate limit(1/s, burst 1)

	time.Sleep(time.Second * 2) // wait for element to be executed
}

输出结果：

$ go run demo.go
>> get element: hello
>> get element: world
>> get element: burst
>> get element: delay
>> get element: limit

如果你不想使用 workqueue.NewRateLimitingQueue(nil) 这个方式来创建队列，还有一个函数偷懒 workqueue.DefaultRateLimitingQueue()，两者是等效的。

注意: 创建一个 RateLimiting Queue 的实例，是允许绑定自定义的 Delaying Queue 的，只要你的 Delaying Queue 实现了 Interface 接口，就可以绑定到 RateLimiting Queue 上。

// 创建一个 RateLimitingQueue 实例
// Create a new RateLimitingQueue config
func NewRateLimitingQueue(conf *RateLimitingQConfig) *RateLimitingQ {
	conf = isRateLimitingQConfigValid(conf)
	conf.DelayingQConfig.cb = conf.cb
	return NewRateLimitingQueueWithCustomQueue(conf, NewDelayingQueue(&conf.DelayingQConfig))
}

RateLimiting Queue 代码解析

Limiter Interface

只要你想要开发一个限速器，实现如下的所有接口，就可以让 RateLimiting Queue 使用你定义的限速逻辑，是不是很开放？哈哈!!!

具体接口定义如下：

// 定义一个限速器接口
// Defines the rate limiter interface.
type RateLimiter interface {
	// 获取一个元素应该等待多长时间
	// When gets an element and gets to decide how long that element should wait
	When(element any) time.Duration

	// 停止追踪这个元素，不让他继续在队列中等待。
	// Forget indicates that an element is finished being retried.  Doesn't matter whether it's for failing
	// or for success, we'll stop tracking it
	Forget(element any)

	// 返回一个元素被限速的次数
	// NumLimitTimes returns back limit times the element has had
	NumLimitTimes(element any) int

	// 关闭限速器
	// Stop stops the limiter
	Stop()
}

TokenBucket Limiter

默认的 TokenBucket 的方式在拿到了令牌只有两种处理方式，"使用"或者"丢弃"令牌，不用把令牌还回去。由于采用的每次拿一个令牌的方式，所以不需要对令牌和处理数量挂钩，同时"丢弃"令牌不需要对令牌和 TokenBucket 做额外处理，所以基于 TokenBucket 的 Limiter 实现非常简单。

// 实现一个基于令牌桶 RateLimiter
// Implements a rate limiter that uses a token bucket.
type BucketRateLimiter struct {
	*rate.Limiter
}

// 创建一个基于令牌桶 RateLimiter
// NewBucketRateLimiter returns a new instance of a bucket rate limiter.
func NewBucketRateLimiter(r float64, b int) *BucketRateLimiter {
	return &BucketRateLimiter{rate.NewLimiter(rate.Limit(r), b)}
}

// 返回一个元素被限速的次数
// Return the number of times an element is limited
func (r *BucketRateLimiter) NumLimitTimes(element any) int {
	// do nothing
	return 0
}

// 停止追踪这个元素，不让他继续在队列中等待。
// Forget indicates that an element is finished being retried.  Doesn't matter whether it's for failing
func (r *BucketRateLimiter) Forget(element any) {
	// do nothing
}

// 关闭令牌桶
// Stop stops the limiter
func (r *BucketRateLimiter) Stop() {
	// do nothing
}

当然那我还实现了一个基于指数退避的 RateLimiter，具体可以参考项目的中 ratelimiter.go 文件。

AddLimited

AddLimited 实际就是调用了 Delaying Queue 的 AddAfter 方法，只不过是将 Limiter 的延迟时间作为 Delaying Queue 的延迟时间。

// 添加元素到队列, 加入到等待队列, 如果有 token 则直接加入到队列
// Add an element to the queue, add it to the waiting queue, and add it to the queue directly if there is a token
func (q *RateLimitingQ) AddLimited(element any) error {
	if q.IsClosed() {
		return ErrorQueueClosed
	}

	err := q.AddAfter(element, q.limiter.When(element)) // 加入到等待队列, add it to the waiting queue
	q.config.cb.OnAddLimited(element)
	return err
}

总结

谢谢你，到此你已经看完 WorkQueue 项目全部的文章了。通过本篇文章向你介绍了 RateLimiting Queue 的定义、使用方法和实现原理。希望能够让你对 RateLimiting Queue 有一个清晰的认识。

RateLimiting Queue 是 WorkQueue 中的核心模块中最高级的 Queue，也是最复杂的一个 Queue。RateLimiting Queue 的实现是基于 Delaying Queue 的，Delaying Queue 的实现是基于 Queue 的。所以说 RateLimiting Queue 包含了 Queue 和 Delaying Queue 的所有特性，同时还有 RateLimiting 的特性。RateLimiting Queue 的实现非常简单，但是使用起来非常的方便，只需要调用 AddLimited 方法就可以了，当然别忘了还有 Forget 和 NumLimitTimes 方法，可以帮你更加灵活控制 RateLimiting Queue 的行为。

当然 RateLimiting Queue 在创建实例的时候，到底选用 Queue 还是 Simple Queue，请根据实际业务情况来决定。