使用Cloudflare Durable Objects构建一个分布式限流器

import { DurableObject } from "cloudflare:workers";
import type { Duration } from "@/lib/duration";
import { ms } from "@/lib/duration";

// 定义存储在 DO 中的状态接口
interface BucketState {
  tokens: number; // 当前剩余令牌数
  lastRefill: number; // 上次填充令牌的时间戳
}

// 限流配置项
export type RateLimitOptions = {
  capacity: number; // 桶的容量（最大令牌数）
  interval: Duration; // 填充周期
  cost?: number; // 本次请求消耗的令牌数，默认为 1
};

export class RateLimiter extends DurableObject {
  private state: BucketState;

  constructor(ctx: DurableObjectState, env: Env) {
    super(ctx, env);
    // 初始化内存状态
    this.state = {
      tokens: 0,
      lastRefill: 0,
    };

    // 使用 blockConcurrencyWhile 确保在处理请求前，先从存储中加载状态
    // 这是保证数据一致性的关键
    ctx.blockConcurrencyWhile(async () => {
      const stored = await ctx.storage.get<BucketState>("bucket");
      if (stored) {
        this.state = stored;
      }
    });
  }

  // 检查是否允许请求的核心逻辑 (RPC 方法)
  checkLimit({ capacity, interval, cost = 1 }: RateLimitOptions): {
    allowed: boolean;
    remaining: number;
    retryAfterMs: number;
  } {
    // 边界检查：如果单次消耗超过总容量，直接拒绝
    if (cost > capacity) {
      return {
        allowed: false,
        remaining: 0,
        retryAfterMs: -1,
      };
    }

    const now = Date.now();
    const intervalMs = ms(interval); // 将时长转换为毫秒
    const rate = capacity / intervalMs; // 计算填充速率（令牌/毫秒）

    // 首次运行时的初始化
    if (this.state.lastRefill === 0) {
      this.state.lastRefill = now;
      this.state.tokens = capacity;
    }

    // 1. 计算自上次填充以来流逝的时间
    const timeSinceLastRefill = now - this.state.lastRefill;

    // 2. 根据速率计算应补充的令牌数
    const tokensToAdd = timeSinceLastRefill * rate;

    // 3. 更新当前令牌数，但不能超过桶的容量
    const currentTokens = Math.min(capacity, this.state.tokens + tokensToAdd);

    // 4. 检查是否有足够的令牌支付本次消耗
    if (currentTokens < cost) {
      // 令牌不足，拒绝请求
      return {
        allowed: false,
        remaining: Math.floor(currentTokens),
        // 计算还需要等待多久才有足够的令牌
        retryAfterMs: Math.ceil((cost - currentTokens) / rate),
      };
    }

    // 5. 令牌充足，扣除消耗并更新状态
    this.state.tokens = currentTokens - cost;
    this.state.lastRefill = now;

    // 异步将新状态写入持久化存储，不阻塞当前请求的返回
    this.ctx.waitUntil(this.ctx.storage.put("bucket", this.state));

    return {
      allowed: true,
      remaining: Math.floor(this.state.tokens),
      retryAfterMs: 0,
    };
  }
}

export default {
  async fetch(request: Request, env: Env) {
    const userId = "user-123"; // 实际场景中应从 Auth Header 获取

    // 1. 根据 userId 获取唯一的 Durable Object ID
    // idFromName 保证同一个名字总是返回相同的 ID
    const id = env.RATE_LIMITER.idFromName(userId);

    // 2. 获取该 ID 对应的 Stub (客户端存根)
    const rateLimiter = env.RATE_LIMITER.get(id);

    // 3. 像调用本地函数一样调用 DO 中的方法 (RPC)
    const result = await rateLimiter.checkLimit({
      capacity: 60, // 一分钟最多 60 次
      interval: "1m",
      // key 参数在这里其实是不需要的，因为 DO 实例本身就是按用户隔离的
    });

    if (!result.allowed) {
      return Response.json(
        {
          message: "Too Many Requests",
          retryAfterSeconds: Math.ceil(result.retryAfterMs / 1000),
        },
        {
          status: 429,
          headers: {
            "Retry-After": String(Math.ceil(result.retryAfterMs / 1000)),
          },
        },
      );
    }

    // 业务逻辑...
    // doSomething();

    return new Response("OK");
  },
};