Goでゼロからつくるレートリミッター:Token BucketアルゴリズムによるAPI保護

Programming tutorial - IT technology blog
Programming tutorial - IT technology blog

APIがサンドバッグになるとき

これは一度や二度ではなく経験したことがある。サービスが本番稼働し、一見問題なく動いているように見えるとき、突然トラフィックが急増するか、誰かがエンドポイントを連打し始める。すると、レスポンスタイムが跳ね上がり、データベースが悲鳴を上げ、ユーザーにはエラーが表示され始める。最悪なのは、これが完全に防ぎ得るものだということだ。

今では、どんな本番APIでも最初に実装するのがレートリミッターだ。後付けでも、理解しないままライブラリを組み込むのでもなく、ゼロから実装できるくらい深く理解しているものとして。Token Bucketアルゴリズムが私の定番だ。シンプルで柔軟性があり、固定ウィンドウ方式よりも実際のトラフィックパターンをはるかにうまく処理できる。

GoでAPIを構築していてまだレートリミッターを追加していないなら、このチュートリアルでゼロから構築する方法を解説する。約100行のコードで、内容をしっかり理解できるはずだ。

保護されていないAPIの問題

レートリミッターがなければ、設定ミスの1クライアントだけでサービスをダウンさせることができる。sleep呼び出しが抜けたスクリプトでそれが起きるのを見たことがある。本物の攻撃ですらなく、誰かの自動化スクリプトがwebhookをタイトなループで叩き続けていただけだ。サーバーはリクエストをキューに積み始め、メモリは上昇し、1分以内に正規のユーザーがタイムアウトし始めた。

レートリミッターは3つの問題を解決する:

  • DoS対策 — 1つのIPがサーバーリソースを独占できない
  • 公平な利用 — ヘビーユーザー1人が他の全員のサービス品質を低下させられない
  • コスト制御 — 下流サービス(データベース、外部API)が大量リクエストで溢れない

Token Bucketアルゴリズムを理解する

トークンを保持するバケツをイメージしてほしい。入ってくるリクエストは毎回1トークンを消費する。トークンは一定のレートで補充される。リクエストが来たときにバケツが空であれば、そのリクエストは拒否される。

2つのパラメータがすべてを定義する:

  • キャパシティ — バケツが保持できる最大トークン数(バースト量を制御)
  • 補充レート — 1秒あたりに追加されるトークン数(持続スループットを制御)

Token Bucketは固定ウィンドウカウント方式に比べて明確な優位性がある。固定ウィンドウは「1分あたり100リクエスト」と言うが、クライアントが最初の1秒で100リクエスト全部を送ることを止める手段がない。Token Bucketは正当なバーストを適切に処理しながら、持続的なレートリミットを適用できる。しばらく待機した後に10リクエストを素早く送るユーザーは通過できるが、ループで1000回APIを叩く人はスロットリングされる。

ライブラリを使わずに自分で作る理由

golang.org/x/time/rateを使う手もある。堅牢で実戦検証済みだ。しかし、内部の仕組みを理解することに意味がある。本番環境で午前3時に何かが壊れたとき、「ライブラリを使いました」ではデバッグの助けにならない。自分で実装すれば30分で終わり、それ以降は動作で迷うことがなくなる。かつて本番のレートリミッターのバグを追跡したら、微妙なクロックの問題だった。実装を知っていたおかげで、ライブラリのソースを一晩読む代わりに10分で修正できた。

GoでRate Limiterを構築する

コアデータ構造

バケツの構造体から始めよう:

package ratelimiter

import (
    "sync"
    "time"
)

type TokenBucket struct {
    capacity   float64
    tokens     float64
    refillRate float64 // 1秒あたりのトークン数
    lastRefill time.Time
    lastAccess time.Time
    mu         sync.Mutex
}

func NewTokenBucket(capacity float64, refillRate float64) *TokenBucket {
    now := time.Now()
    return &TokenBucket{
        capacity:   capacity,
        tokens:     capacity, // 最初は満杯
        refillRate: refillRate,
        lastRefill: now,
        lastAccess: now,
    }
}

トークンにintegerではなくfloat64を使っている。これにより補充計算がシンプルになる。リクエスト間で小数点以下でトークンが蓄積されるときに丸め誤差が生じない。

Allow()メソッド

毎回の呼び出しで2つの操作が行われる。経過時間に基づいてトークンを補充し、次に1トークンを消費しようとする:

func (tb *TokenBucket) Allow() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()

    now := time.Now()
    elapsed := now.Sub(tb.lastRefill).Seconds()

    // 経過時間に基づいてトークンを追加
    tb.tokens += elapsed * tb.refillRate
    if tb.tokens > tb.capacity {
        tb.tokens = tb.capacity
    }
    tb.lastRefill = now
    tb.lastAccess = now

    // トークンの消費を試みる
    if tb.tokens >= 1 {
        tb.tokens--
        return true
    }
    return false
}

mutexは必須だ。GoのHTTPサーバーは各リクエストを独自のgoroutineで処理するため、同期なしではレースコンディションが発生し、トークン数が破損する。初期バージョンでこれを経験した。ストレステストでは同時負荷下での不規則な動作が見られたが、単独では再現がほぼ不可能だった。

クライアントごとのレート制限

単一のグローバルバケツではAPI全体を一度にスロットリングしてしまい、それは望む動作ではない。クライアントごとに1つのバケツが必要だ。IPアドレスやAPIキーをキーとして使う:

type RateLimiter struct {
    clients    map[string]*TokenBucket
    mu         sync.RWMutex
    capacity   float64
    refillRate float64
}

func NewRateLimiter(capacity, refillRate float64) *RateLimiter {
    rl := &RateLimiter{
        clients:    make(map[string]*TokenBucket),
        capacity:   capacity,
        refillRate: refillRate,
    }
    go rl.cleanup()
    return rl
}

func (rl *RateLimiter) getBucket(clientID string) *TokenBucket {
    rl.mu.RLock()
    bucket, exists := rl.clients[clientID]
    rl.mu.RUnlock()

    if exists {
        return bucket
    }

    rl.mu.Lock()
    defer rl.mu.Unlock()
    // 書き込みロック取得後に二重確認
    if bucket, exists = rl.clients[clientID]; exists {
        return bucket
    }
    bucket = NewTokenBucket(rl.capacity, rl.refillRate)
    rl.clients[clientID] = bucket
    return bucket
}

func (rl *RateLimiter) Allow(clientID string) bool {
    return rl.getBucket(clientID).Allow()
}

getBucket()内のダブルチェックパターンに注目してほしい。読み取りロックを解放してから書き込みロックを取得するまでの間に、別のgoroutineが同じバケツを作成してしまう可能性がある。これを省くと、高トラフィック下で断続的な重複アロケーションが発生する。並行アクセス下での遅延初期化におけるGoの標準的なアプローチだ。

古くなったバケツのクリーンアップ

放置すると、クライアントマップはユニークな呼び出し元ごとに1エントリずつ増え続ける。数千のユニークIPを持つ公開APIではメモリが際限なく増加する。バックグラウンドgoroutineが削除を処理する:

func (rl *RateLimiter) cleanup() {
    ticker := time.NewTicker(5 * time.Minute)
    for range ticker.C {
        rl.mu.Lock()
        for id, bucket := range rl.clients {
            bucket.mu.Lock()
            if time.Since(bucket.lastAccess) > 10*time.Minute {
                delete(rl.clients, id)
            }
            bucket.mu.Unlock()
        }
        rl.mu.Unlock()
    }
}

HTTPミドルウェア

ミドルウェアチェーンに組み込むだけで完成だ:

func RateLimitMiddleware(rl *RateLimiter) func(http.Handler) http.Handler {
    return func(next http.Handler) http.Handler {
        return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
            // プロキシの背後の本番環境ではX-Forwarded-Forを使用する
            clientIP := r.RemoteAddr

            if !rl.Allow(clientIP) {
                w.Header().Set("Retry-After", "1")
                w.Header().Set("X-RateLimit-Limit", "10")
                http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
                return
            }
            next.ServeHTTP(w, r)
        })
    }
}

func main() {
    // バースト10トークン、持続レート2トークン/秒
    limiter := NewRateLimiter(10, 2)

    mux := http.NewServeMux()
    mux.HandleFunc("/api/data", handleData)

    handler := RateLimitMiddleware(limiter)(mux)
    http.ListenAndServe(":8080", handler)
}

バーストとスロットルの動作をテストする

テストを省略しないこと。バーストの動作を事前に確認しておくと、本番環境での当てずっぽうのチューニングを避けられる:

func TestTokenBucket_BurstThenThrottle(t *testing.T) {
    // キャパシティ5トークン、補充レート1トークン/秒
    bucket := NewTokenBucket(5, 1)

    // バースト: 5回の連続リクエストはすべて通過するはず
    for i := 0; i < 5; i++ {
        if !bucket.Allow() {
            t.Fatalf("リクエスト %d は許可されるはずです", i+1)
        }
    }

    // 6番目は拒否されるはず — バケツが空
    if bucket.Allow() {
        t.Fatal("6番目のリクエストは拒否されるはずです")
    }

    // 1トークンが補充されるまで待機
    time.Sleep(1100 * time.Millisecond)
    if !bucket.Allow() {
        t.Fatal("補充後はリクエストが許可されるはずです")
    }
}

チューニングと本番環境のヒント

適切なパラメータの選び方

数値を間違えることが最もよくあるミスだ。公開APIに対する私の出発点:

  • キャパシティは想定される正当なバーストの2〜3倍に設定する(ユーザーはエラー後に素早くリトライすることがある)
  • 補充レートは持続的なリクエスト/秒の予算に設定する
  • 認証済みエンドポイント: より寛大に — キャパシティ50トークン、補充10/秒
  • 未認証または機密性の高いエンドポイント: 厳格に — キャパシティ5トークン、補充1/秒

分散レートリミッター

インメモリ実装は単一サーバーでは完璧に機能する。水平スケールすると問題がすぐに現れる。各サーバーが独自のバケツ状態を持つため、クライアントは5台のサーバーそれぞれに10 req/秒でアクセスして合計50 req/秒を達成できてしまう。分散環境では、トークン状態をアトミックなLuaスクリプトを使ってRedisに移す。アルゴリズムはそのまま移行できる。同じ計算、異なるストレージバックエンドだ。シングルノードを卒業したら、別途深掘りする価値がある。

ヘッダーでクライアントに状況を伝える

コンテキストのない429レスポンスはデバッグが辛い。常に返すようにすること:

  • X-RateLimit-Limit — バケツのキャパシティ
  • X-RateLimit-Remaining — 現在のトークン数(消費前に読み取る)
  • Retry-After — 次のトークンが利用可能になるまでの秒数

行儀の良いクライアントはこれらのヘッダーを見て自動的にバックオフする。行儀の悪いクライアントでも、少なくともログで指摘できるものが残る。

次のステップ

上記の実装は約100行のGoで、並行アクセスを正しく処理し、古い状態を自動的にクリーンアップし、標準的なHTTPミドルウェアに組み込める。一気に読めるほどシンプルで、適用が必要なとき、午前3時にデバッグするとき、または次のエンジニアに一行ずつ説明するときに役立つ。

レートリミッターは障害報告書ではなく、最初のコミットに含まれるべきものだ。ここで紹介したデフォルト値から始め、実際のエンドポイントに対してロードテストを実行し、推測ではなく実際のトラフィック数値からチューニングしよう。

Share: