golang令牌桶实现 [go-rate] 速率限制器-木庄网络博客

本文摘自网络，作者，侵删。

go-rate是速率限制器库,基于 Token Bucket(令牌桶)算法实现。 go-rate被用在LangTrend的生产中用于遵守GitHub API速率限制。

速率限制可以完成一些特殊的功能需求，包括但不限于服务器端垃圾邮件保护、防止api调用饱和等。

库使用说明

我们首先构造一个限流器对象：

limiter := NewLimiter(10, 1);

这里有两个参数：

上述的限流器的含义是：拥有一个容量为1的令牌桶，以每钞10个的速度向桶中放令牌。

除了直接指定每秒产生的 Token 个数外，还可以用 Every 方法来指定向 Token 桶中放置 Token 的间隔，例如：

limiter := NewLimiter(Every(100 * time.Millisecond), 1);

以上就表示每 100ms 往桶中放一个 Token。本质上也就是一秒钟产生 10 个。

Limiter 提供了三类方法供用户消费 Token，用户可以每次消费一个 Token，也可以一次性消费多个 Token。
而每种方法代表了当 Token 不足时，各自不同的对应手段。

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

Wait 实际上就是 WaitN(ctx,1)。

当使用 Wait 方法消费 Token 时，如果此时桶内 Token 数组不足 (小于 N)，那么 Wait 方法将会阻塞一段时间，直至 Token 满足条件。如果充足则直接返回。

这里可以看到，Wait 方法有一个 context 参数。我们可以设置 context 的 Deadline 或者 Timeout，来决定此次 Wait 的最长时间。

Allow 实际上就是 AllowN(time.Now(),1)。

AllowN 方法表示，截止到某一时刻，目前桶中数目是否至少为 n 个，满足则返回 true，同时从桶中消费 n 个 token。
反之返回不消费 Token，false。

通常对应这样的线上场景，如果请求速率过快，就直接丢到某些请求。

Reserve 相当于 ReserveN(time.Now(), 1)。

ReserveN 的用法就相对来说复杂一些，当调用完成后，无论 Token 是否充足，都会返回一个 Reservation * 对象。

你可以调用该对象的 Delay() 方法，该方法返回了需要等待的时间。如果等待时间为 0，则说明不用等待。必须等到等待时间之后，才能进行接下来的工作。

或者，如果不想等待，可以调用 Cancel() 方法，该方法会将 Token 归还。

￥0元 老貘

一个与时俱进的Go编程知识库。