[Go] golang.org/x/time/rate でレイトリミット

レイトリミットを実装するために golang.org/x/time/rate を採用したものの、しばらく経ってから見返すとその仕様を完全に忘れていてショックだったので、こちらに記事として備忘録的に要点をまとめていきたいと思います。

トークンバケット

golang.org/x/time/rate は「トークンバケット」というアルゴリズムを Go で実装したパッケージになります。
このパッケージは一見すると使い方が分かりづらいのですが、「トークンバケット」を理解すると合点がいくデザインとなっているので、まずはこちらを理解しましょう。

一応 Wikipedia にもページがあるのですが（トークンバケット - Wikipedia）、golang.org/x/time/rate パッケージのドキュメントを読むのが端的で分かりやすいと個人的には思いました。Limiter 構造体に対する説明として記載されておりまして、以下のリンク先から読むことができます。

• rate · pkg.go.dev

使い方

例えば、こんなコードがあるとします：

package main

import (
    "fmt"
)

func main() {
    for i := 0; i < 10; i++ {
        fmt.Printf("%d\n", i)
    }
}

コードを見れば自明ですが、このコードを実行すると一瞬で0から9までが出力されます。

0
1
2
3
4
5
6
7
8
9

Program exited.

それでは次に、レイトリミットを入れていきます。
ポイントは NewLimiter 関数と Wait 関数です：

package main

import (
    "context"
    "fmt"
    "golang.org/x/time/rate"
)

func main() {
    l := rate.NewLimiter(5.0, 1)

    ctx := context.Background()

    for i := 0; i < 10; i++ {
        if err := l.Wait(ctx); err != nil {
            panic(err)
        }

        fmt.Printf("%d\n", i)
    }
}

上のプログラムを実行すると、各数字が一定間隔で出力されるようになり、だいたい2秒経過して処理を終えると思います。Wait 関数はその名の通り「待つ」関数です。一定間隔で出力されているのはこの関数が制御を一定時間ブロックしているためです。そして、どの程度待つかをしているのは NewLimiter 関数の引数です：

func NewLimiter(r Limit, b int) *Limiter

r がレートにあたる値で、1秒間に最大何回実行するかを決める数になります。r は Limit 型ですが、これは構造体などではなく単なる float64 なので、float64 実値で指定すれば良いと思います。上のプログラムでは、1秒間に5回実行されるように設定していましたので、10回の出力を終えるまでに2秒程度かかった、ということになります。ちなみに b はバーストと呼ばれる値で、何回までは連続して実行しても良いか、を指定するものです。バーストのユースケースについては「Q&A」の節で考察しておりますが、基本的には1を指定するで良いと思います。

Limit は実値で指定する以外に、Every という便利関数を使う方法もあります：

func Every(interval time.Duration) Limit

Every 関数を使うこと「1秒間に何回」という指定の代わりに、「実行毎にどの程度の間隔をあけるか」という形で指定することができます。先程のプログラム内の NewLimiter 関数呼び出しは、以下のように実装しても等価になります（1秒 / 5回 = 0.2秒 なので）：

        interval := rate.Every(time.Millisecond * 200)
    l := rate.NewLimiter(interval, 1)

Q&A

Every 関数をあえて使う必要がある？

Every 関数はただのユーティリティ関数なので、無理に使う必要はありません。プログラムによっては Every 関a数の指定形式の方が分かりやすく書ける時もあると思いますので、そういった場合にのみ使えば良いです。

（ちなみに、どうしてか以下のような冗長なサンプルコードが蔓延しているので、あえてこういった Q&A を設けさせていただいております。🤔）

    interval := rate.Every(time.Second / time.Duration(5))
    l := rate.NewLimiter(interval, 1)

バーストを1以外に指定するユースケースがわからない

レイトリミットを設定しつつAPIを連続してコールするプログラムがあるとします。APIサーバが待ち受けているコネクション数がnで、n個のコールは並列して呼んで良いというときに、バーストをnと設定するといい感じにn個まではウェイトすることなくコールしてくれるので便利、ということだと思います。

Allow メソッドや Reserve メソッドはどう使うの？

私のコードでは必要なかったのでわからないですが…用意されているということは想定されるユースケースがあるということだと思うので、機会があれば使っていきたいですね。その際にはこちらの記事に解説を追記したいと思います。

"go.uber.org/ratelimit" はどう？

検索すると、go.uber.org/ratelimit に限らず Go のレイトリミット系のパッケージはたくさん見つかりますね。
各パッケージの機能を比較した訳ではないのですが、プレフィクスが golang.org/x となっている通り、Go の準標準ライブラリといった位置付けのパッケージですので、個人的にはこちらを好んで使っています。（機能的に優れているなど、他にオススメのパッケージがあれば教えていただけると助かります）