この記事は独立した実装解説です。リアルタイム面接支援という利用背景は、こちらの設計背景を参照してください。

TypeScriptでToken Bucketを実装する：ユーザー単位のレート制限を公平にする

短時間に集中するリクエストを即座に全部落とすと、実際に操作しているユーザーまで止めてしまいます。Token Bucket は、平均レートを守りながら、容量の範囲で一時的なバーストを許す方式です。

この記事では、次の境界を明示した実装を作ります。

• バーストは capacity 回まで許可する
• 経過時間に応じてトークンを連続的に補充する
• 拒否時にも、次に試せる時刻を retryAfterMs として返す
• プロセスが複数になったら、Redis Lua で「補充と消費」を一つの原子的操作にする

先に決めるべき単位

レート制限で最初に決めるのはアルゴリズムではなくキーです。IP だけでは NAT 配下の正規ユーザーを巻き込み、ユーザー ID だけでは未認証の入口を守れません。

たとえば、ログイン済み API なら rate:user:<userId>:answer のように「ユーザー × 操作種別」で分けます。音声の小さなチャンクと、重い回答生成を同じバケットに入れるのも避けます。コストが違う操作は、別バケットか重み付きの cost で扱います。

ここでは capacity = 3、refillPerSecond = 2 とします。空のバケットなら、1 トークンを得るまで 500 ms 待つ計算です。

単一プロセスでの最小実装

状態は「残りトークン」と「最後に補充を計算した時刻」だけです。現在時刻までの経過時間から補充し、容量を超えないように丸めてから消費します。

type Clock = () => number;

type TakeResult = {
  allowed: boolean;
  remaining: number;
  retryAfterMs: number;
};

export class TokenBucket {
  #tokens: number;
  #lastMs: number;

  constructor(
    private readonly capacity: number,
    private readonly refillPerSecond: number,
    private readonly now: Clock = Date.now,
  ) {
    if (!Number.isFinite(capacity) || capacity <= 0) {
      throw new RangeError("capacity must be positive");
    }
    if (!Number.isFinite(refillPerSecond) || refillPerSecond <= 0) {
      throw new RangeError("refillPerSecond must be positive");
    }

    this.#tokens = capacity;
    this.#lastMs = now();
  }

  tryTake(cost = 1): TakeResult {
    if (!Number.isFinite(cost) || cost <= 0 || cost > this.capacity) {
      throw new RangeError("cost must be in (0, capacity]");
    }

    const now = this.now();
    const elapsedMs = Math.max(0, now - this.#lastMs);

    this.#tokens = Math.min(
      this.capacity,
      this.#tokens + (elapsedMs * this.refillPerSecond) / 1_000,
    );
    // 壁時計が戻っても、未来のトークンを作らない
    this.#lastMs = Math.max(this.#lastMs, now);

    if (this.#tokens >= cost) {
      this.#tokens -= cost;
      return { allowed: true, remaining: this.#tokens, retryAfterMs: 0 };
    }

    return {
      allowed: false,
      remaining: this.#tokens,
      retryAfterMs: Math.ceil(
        ((cost - this.#tokens) * 1_000) / this.refillPerSecond,
      ),
    };
  }
}

重要なのは、拒否したときも補充後の tokens と lastMs を保存することです。拒否を「何も変えない」と扱うと、部分的に補充された量が毎回消え、返した Retry-After と実際の挙動がずれます。

境界をテストで固定する

Date.now() を直接呼ぶと時間のテストが不安定になります。クロックを注入すると、待機なしでバースト、部分補充、時計の逆行を検証できます。

import { expect, test } from "bun:test";

test("capacity までのバーストを許可し、次の待ち時間を返す", () => {
  let ms = 0;
  const bucket = new TokenBucket(3, 2, () => ms);

  expect([
    bucket.tryTake().allowed,
    bucket.tryTake().allowed,
    bucket.tryTake().allowed,
  ]).toEqual([true, true, true]);

  expect(bucket.tryTake()).toEqual({
    allowed: false,
    remaining: 0,
    retryAfterMs: 500,
  });
});

test("部分補充し、capacity を超えない", () => {
  let ms = 0;
  const bucket = new TokenBucket(5, 2, () => ms);
  bucket.tryTake(5);

  ms = 750;
  expect(bucket.tryTake(1)).toEqual({
    allowed: true,
    remaining: 0.5,
    retryAfterMs: 0,
  });

  ms = 10_000;
  expect(bucket.tryTake(5)).toEqual({
    allowed: true,
    remaining: 0,
    retryAfterMs: 0,
  });
});

test("壁時計が戻ってもトークンを増やさない", () => {
  let ms = 1_000;
  const bucket = new TokenBucket(1, 1, () => ms);
  bucket.tryTake();

  ms = 500;
  expect(bucket.tryTake()).toEqual({
    allowed: false,
    remaining: 0,
    retryAfterMs: 1_000,
  });
});

この実装では Bun で 4 テスト、6 assertions を実行しました。実サービスでは、同じリクエストを複数回実行しない設計も別途必要です。レート制限は重複実行を防ぐ仕組みではありません。

複数プロセスでは Redis に状態を寄せる

メモリ内の TokenBucket は 1 プロセス内でしか正しく動きません。ワーカーが 3 台あれば、各ワーカーが独立してバーストを許してしまいます。

Redis の EVAL なら、読み取り・補充・消費・TTL 更新を一つのスクリプトとして実行できます。拒否時にも状態を保存する点は、先ほどの TypeScript 実装と同じです。

-- KEYS[1] = rate:user:<id>:answer
-- ARGV = capacity, refillPerSecond, nowMs, cost
local capacity = tonumber(ARGV[1])
local refillPerMs = tonumber(ARGV[2]) / 1000
local nowMs = tonumber(ARGV[3])
local cost = tonumber(ARGV[4])

if cost <= 0 or cost > capacity then
  return redis.error_reply("invalid cost")
end

local saved = redis.call("HMGET", KEYS[1], "tokens", "at")
local tokens = tonumber(saved[1]) or capacity
local at = tonumber(saved[2]) or nowMs
local elapsedMs = math.max(0, nowMs - at)

tokens = math.min(capacity, tokens + elapsedMs * refillPerMs)
at = math.max(at, nowMs)

local ttlMs = math.max(1, math.ceil(capacity / refillPerMs))
if tokens < cost then
  local retryAfterMs = math.ceil((cost - tokens) / refillPerMs)
  redis.call("HSET", KEYS[1], "tokens", tokens, "at", at)
  redis.call("PEXPIRE", KEYS[1], ttlMs)
  return {0, tostring(tokens), retryAfterMs}
end

tokens = tokens - cost
redis.call("HSET", KEYS[1], "tokens", tokens, "at", at)
redis.call("PEXPIRE", KEYS[1], ttlMs)
return {1, tostring(tokens), 0}

HTTP の応答まで決める

拒否時に単に 429 を返すだけでは、クライアント側は適切に待てません。上の結果を HTTP に写すなら、少なくとも Retry-After を付けます。

const result = await takeFromRedis({
  key: `rate:user:${userId}:answer`,
  capacity: 3,
  refillPerSecond: 2,
  cost: 1,
});

if (!result.allowed) {
  return new Response("Too Many Requests", {
    status: 429,
    headers: {
      "Retry-After": String(Math.ceil(result.retryAfterMs / 1_000)),
      "Cache-Control": "no-store",
    },
  });
}

Retry-After を丸め上げるのは、0 秒と返してすぐ再試行させないためです。一方、ログには key の生値ではなく、操作種別・許可/拒否・待機時間・残量だけを記録すると、ユーザー識別子を不要に広げずに調整できます。

実装前の確認項目

• 1 時間あたりの上限ではなく、通常の操作が短時間にどれだけ連続するかを計測したか
• 操作の重さに応じて cost またはバケットを分けたか
• 同じバケットを消費するすべての入口で、同じ Redis スクリプトを使っているか
• 429 の理由と Retry-After をクライアントに返しているか
• バースト、部分補充、容量上限、時計の逆行、複数ワーカーの競合をテストしたか

Token Bucket の本体は小さいですが、キーの境界と原子的な更新を曖昧にすると、公平性も保護も崩れます。