golang.org/x/time/rateでレートリミットを実装してみた

はじめに

先日とある API 関連の本を読んでいたときに、レートリミットについて触れている箇所がありました。
それを読んでいて、ふと「自分でレートリミットの実装をしたことないなぁ」「Go でどうやってレートリミットを実装するんだろう」と思ったので、この記事では golang.org/x/time/rate を使ってレートリミットを実装してみたことをまとめます。

システムの安定性を保つためには、外部からのリクエスト数を適切に制御することが重要です。この記事が、Go で API サーバーなどを開発している人の参考になれば幸いです。

golang.org/x/time/rate とは

golang.org/x/time/rate は、準標準パッケージの一つで、トークンバケット (Token Bucket) と呼ばれるアルゴリズムに基づいたレートリミッターを提供します。

なぜレートリミットが必要か？

外部に公開している API サーバーなどのシステムを運用する上で、以下のような理由からレートリミットは不可欠です。

• 過負荷からの保護: 特定のユーザーやクライアントからの大量リクエストによるサーバーダウンを防ぎます。
• 公平なリソース分配: 全てのユーザーが安定してサービスを利用できるよう、リソースを公平に分配します。
• 悪意のある攻撃の緩和: ブルートフォース攻撃や DDoS 攻撃のような、短時間に大量のリクエストを送る攻撃を緩和します。

トークンバケットアルゴリズムの仕組み

rateパッケージを理解するために、トークンバケットの概念について説明します。これは「トークンが入ったバケツ」に例えることができます。

1. バケツ (Bucket): リクエストを処理する権利である「トークン」を貯めておく場所です。このバケツの容量が burst (バースト) 量にあたります。
2. トークンの補充 (Refill): バケツには、一定のペースで新しいトークンが補充されます。この補充される速さが rate (レート) です。
3. リクエストの処理 (Consume): リクエストが 1 つ来るたびに、バケツからトークンを 1 つ消費します。

もしリクエストが来たときにバケツにトークンがあれば、即座にリクエストは許可されます。しかし、トークンが空っぽの場合、新しいトークンが補充されるまで待つか、リクエストを諦める（拒否する）しかありません。

この仕組みの優れた点は、平均レートを制御しつつ、一時的なリクエストの集中（バースト）にも対応できることです。

• rate.Limit (レート): 1 秒あたりにバケツに補充されるトークンの数です。これにより、長期的に見た場合の平均リクエストレートが決定されます。例えば rate.Limit(10) であれば、1 秒間に 10 個のペースでトークンが補充されます。
• burst (バースト): バケツの最大容量です。普段リクエストが少なくトークンが貯まっている状態であれば、最大でこの burst 数まで、レートを超えた連続リクエストを即座に処理できます。これにより、アプリケーションの応答性を損なわずに、突発的なトラフィックをさばくことが可能になります。

インストール

まずはパッケージをインストールします。

go get golang.org/x/time/rate

基本的な使い方

rate.NewLimiter を使って、リミッターを生成します。引数には、1 秒あたりに許可するイベント数（rate.Limit）と、バケツの大きさ（burst）を指定します。

import (
	"context"
	"fmt"
	"time"

	"golang.org/x/time/rate"
)

func main() {
	// 1秒あたり2リクエスト、バーストは5まで許容
	limiter := rate.NewLimiter(2, 5)

	ctx := context.Background()

	for i := 0; i < 10; i++ {
		// トークンが取得できるまで待機
		if err := limiter.Wait(ctx); err != nil {
			fmt.Println("error:", err)
			return
		}
		fmt.Printf("%s: リクエスト %d を処理しました
", time.Now().Format("15:04:05.000"), i+1)
	}
}

Wait メソッド

limiter.Wait(ctx) は、リクエストが許可されるまでブロックするメソッドです。リミッターのレートを超えるリクエストが来た場合、トークンが利用可能になるまで処理を待ちます。

上記のコードを実行すると、最初の 5 リクエストはバースト許容量のおかげで即座に処理されますが、その後は 1 秒あたり 2 リクエストのペースに制限されるため、約 0.5 秒ごとに処理が進むのが分かります。

Allow メソッド

limiter.Allow() は、Wait とは異なり、待機せずに現在のリクエストが許可されるかどうかを bool 値で返します。すぐに結果が知りたい場合や、リクエストを待たせるのではなく、拒否したい場合（429 Too Many Requests を返すなど）に便利です。

if limiter.Allow() {
    fmt.Println("リクエストを処理します。")
} else {
    fmt.Println("リクエストが多すぎます。")
}

HTTP サーバーでの実装例

Web サーバーのミドルウェアとしてレートリミットを実装するのが一般的な使い方です。以下に net/http を使った簡単な例を示します。

package main

import (
	"log"
	"net/http"

	"golang.org/x/time/rate"
)

var limiter = rate.NewLimiter(1, 3) // 1秒あたり1リクエスト、バケットは最大3リクエストを保持

func rateLimitMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !limiter.Allow() {
			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
			log.Println("Rate limit exceeded.")
			return
		}
		next.ServeHTTP(w, r)
	})
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Hello, World!"))
	})

	// ミドルウェアを適用
	handler := rateLimitMiddleware(mux)

	log.Println("Server starting on port 8080...")
	if err := http.ListenAndServe(":8080", handler); err != nil {
		log.Fatalf("Could not start server: %s\n", err)
	}
}

このコードでは、グローバルな limiter 変数を作成し、すべてのリクエストに対してこの単一のインスタンスでレートチェックを行っています。具体的には以下の通りです。

1. グローバルリミッター: rate.NewLimiter(1, 3) で、サーバー全体で共有されるリミッターを一つだけ生成します。これは「1 秒あたり 1 リクエストを許可し、最大 3 リクエストまでのバーストを許容する」というルールを意味します。
2. ミドルウェア: rateLimitMiddleware は、受け取ったリクエストを次のハンドラに渡す前に、limiter.Allow() を呼び出します。ここでトークンが取得できなければ、クライアントに 429 Too Many Requests エラーを返して処理を中断します。
3. 適用: rateLimitMiddleware を http.NewServeMux() で作成したハンドラに適用することで、すべてのエンドポイント (/) へのリクエストがこのレートリミットの対象となります。

この方法は、サーバー全体のリクエスト総量を制限したい場合に有効です。例えば、外部 API の呼び出し回数に上限がある場合や、データベースへの負荷を一定以下に保ちたい場合などに利用できます。

この例では、全てのリクエストに対して単一のリミッターを適用しています。
しかし、実際には IP アドレスごとなど、リクエスト元に応じてリミッターを管理したいケースが多いでしょう。その場合は、IP アドレスをキーにしたマップでリミッターを管理するような実装が考えられます。

package main

import (
	"fmt"
	"net"
	"net/http"
	"sync"
	"time"

	"golang.org/x/time/rate"
)

var (
	limiters = make(map[string]*rate.Limiter)
	mu       sync.Mutex
)

// getLimiter は指定されたIPアドレスに対応するレートリミッターを返す
func getLimiter(ip string) *rate.Limiter {
	mu.Lock()
	defer mu.Unlock()

	limiter, exists := limiters[ip]
	if !exists {
		// 2秒ごとに1リクエストを許可し、バケットは最大5リクエストを保持
		limiter = rate.NewLimiter(rate.Every(2*time.Second), 5)
		limiters[ip] = limiter
	}

	return limiter
}

// ミドルウェアとしてレートリミットを適用
func limitMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ip, _, err := net.SplitHostPort(r.RemoteAddr)
		if err != nil {
			http.Error(w, "Internal Server Error", http.StatusInternalServerError)
			return
		}

		limiter := getLimiter(ip)
		if !limiter.Allow() {
			http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
			return
		}

		next.ServeHTTP(w, r)
	})
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, "Hello, Docker with Rate Limiting!")
	})

	// レートリミットミドルウェアを適用
	handlerWithLimit := limitMiddleware(mux)

	fmt.Println("Server starting on port 8080...")
	if err := http.ListenAndServe(":8080", handlerWithLimit); err != nil {
		panic(err)
	}
}

この実装では、以下の流れで IP アドレスに基づいたレートリミットを実現しています。

1. IP アドレスの取得: ミドルウェア内で r.RemoteAddr からリクエスト元の IP アドレスを取得します。net.SplitHostPort を使うことで、IP:Port という形式の文字列から IP アドレス部分だけを安全に抽出できます。
2. リミッターの取得・生成: getLimiter 関数が、IP アドレスをキーとして limiters マップを検索します。
   • もし対応するリミッターが既に存在すれば、それを返します。
   • 存在しない場合は、新しいリミッターを生成してマップに保存し、それを返します。
3. 排他制御: limiters マップへのアクセスは、複数のゴルーチン（リクエスト）から同時に行われる可能性があるため、sync.Mutex を使って競合状態を防いでいます。これにより、マップの読み書きがアトミックに行われ、データが破損するのを防ぎます。
4. レートチェック: 取得したリミッターで Allow() を呼び出し、リクエストを許可するかどうかを判断します。

これにより、クライアントの IP アドレスごとに独立したレート制御が可能になり、特定のユーザーからの大量リクエストが他のユーザーに影響を与えるのを防ぐことができます。

※上記はあくまで簡単な例です。実際には、メモリリークを防ぐために、長時間アクセスがない IP アドレスのエントリを定期的にマップから削除するなどの仕組みが必要になります。

まとめ

golang.org/x/time/rate を使うと、非常にシンプルにレートリミットを実装できることが分かりました。

• NewLimiter で簡単にリミッターを作成できる。
• Wait で待機させるか、Allow で拒否するかを選べる。
• HTTP ミドルウェアとして組み込むことで、API の保護が容易になる。

Go でアプリケーションでレートリミットを実装する際には、便利なパッケージでした。