設計 Rate Limiter — Token Bucket、Sliding Window 與分散式限流實戰 | 資料結構與演算法

/**
 * Token Bucket Rate Limiter
 * - 允許突發流量（burst）
 * - O(1) 時間與空間複雜度
 * - 使用「懶惰補充」策略：只在消耗時才計算應補充的令牌數
 */
class TokenBucket {
  private capacity: number;       // 桶的最大容量
  private tokens: number;         // 當前令牌數
  private refillRate: number;     // 每秒補充速率
  private lastRefillTime: number; // 上次補充時間（ms）

  constructor(capacity: number, refillRate: number) {
    this.capacity = capacity;
    this.refillRate = refillRate;
    this.tokens = capacity;       // 初始滿桶
    this.lastRefillTime = Date.now();
  }

  /** 懶惰補充：只在需要時計算經過時間並補充令牌 */
  private refill(): void {
    const now = Date.now();
    const elapsedSeconds = (now - this.lastRefillTime) / 1000;
    const tokensToAdd = elapsedSeconds * this.refillRate;

    this.tokens = Math.min(this.capacity, this.tokens + tokensToAdd);
    this.lastRefillTime = now;
  }

  /** 嘗試消耗 count 個令牌，回傳是否允許 */
  consume(count: number = 1): boolean {
    this.refill();

    if (this.tokens >= count) {
      this.tokens -= count;
      return true;  // 允許
    }
    return false;   // 拒絕
  }

  getStatus(): { tokens: number; capacity: number } {
    this.refill();
    return { tokens: Math.floor(this.tokens), capacity: this.capacity };
  }
}

/** 管理多用戶的 Token Bucket Rate Limiter */
class TokenBucketRateLimiter {
  private buckets: Map<string, TokenBucket> = new Map();
  private capacity: number;
  private refillRate: number;

  constructor(capacity: number, refillRate: number) {
    this.capacity = capacity;
    this.refillRate = refillRate;
  }

  isAllowed(userId: string): boolean {
    if (!this.buckets.has(userId)) {
      this.buckets.set(
        userId,
        new TokenBucket(this.capacity, this.refillRate)
      );
    }
    return this.buckets.get(userId)!.consume();
  }

  getStatus(userId: string): { tokens: number; capacity: number } | null {
    return this.buckets.get(userId)?.getStatus() ?? null;
  }

  /** 定期清理長時間未活躍的桶，防止記憶體洩漏 */
  cleanup(ttlMs: number = 60_000): void {
    console.log(`Cleanup: ${this.buckets.size} buckets in memory`);
  }
}

/**
 * Sliding Window Counter Rate Limiter
 * - 精確度接近 Sliding Window Log
 * - 記憶體效率如同 Fixed Window Counter
 * - O(1) 時間與空間複雜度
 */
class SlidingWindowCounter {
  private windowSize: number;  // 窗口大小（秒）
  private limit: number;       // 窗口內最大請求數
  private counters: Map<string, number> = new Map();

  constructor(windowSizeSeconds: number, limit: number) {
    this.windowSize = windowSizeSeconds;
    this.limit = limit;
  }

  private getCurrentWindowIndex(): number {
    return Math.floor(Date.now() / 1000 / this.windowSize);
  }

  /** 加權估算當前滑動窗口內的請求數 */
  private getEstimatedCount(userId: string): number {
    const now = Date.now() / 1000;
    const currentWindowIndex = Math.floor(now / this.windowSize);
    const previousWindowIndex = currentWindowIndex - 1;

    // 當前窗口已過的比例
    const currentWindowProgress =
      (now % this.windowSize) / this.windowSize;

    const currentCount =
      this.counters.get(`${userId}:${currentWindowIndex}`) ?? 0;
    const previousCount =
      this.counters.get(`${userId}:${previousWindowIndex}`) ?? 0;

    // 加權：前窗口貢獻 (1 - progress)，當前窗口全部計入
    return previousCount * (1 - currentWindowProgress) + currentCount;
  }

  isAllowed(userId: string): boolean {
    const estimatedCount = this.getEstimatedCount(userId);

    if (estimatedCount >= this.limit) {
      return false; // 限流
    }

    // 增加當前窗口計數
    const currentWindowIndex = this.getCurrentWindowIndex();
    const key = `${userId}:${currentWindowIndex}`;
    this.counters.set(key, (this.counters.get(key) ?? 0) + 1);

    return true; // 允許
  }

  getRemainingRequests(userId: string): number {
    const estimated = this.getEstimatedCount(userId);
    return Math.max(0, this.limit - Math.ceil(estimated));
  }

  getResetTime(): number {
    const currentWindowIndex = this.getCurrentWindowIndex();
    return (currentWindowIndex + 1) * this.windowSize;
  }

  /** 清理舊窗口資料，只保留最近兩個窗口 */
  cleanup(): void {
    const currentWindowIndex = this.getCurrentWindowIndex();
    const validIndexes = new Set([
      `${currentWindowIndex}`,
      `${currentWindowIndex - 1}`,
    ]);

    for (const key of this.counters.keys()) {
      const windowIndex = key.split(":").pop()!;
      if (!validIndexes.has(windowIndex)) {
        this.counters.delete(key);
      }
    }
  }
}

/**
 * Fixed Window Counter Rate Limiter
 * - 最簡單的限流演算法
 * - 有邊界問題但適合粗粒度限流
 */
class FixedWindowCounter {
  private windowSize: number;
  private limit: number;
  private counters: Map<string, { count: number; windowIndex: number }> =
    new Map();

  constructor(windowSizeSeconds: number, limit: number) {
    this.windowSize = windowSizeSeconds;
    this.limit = limit;
  }

  private getCurrentWindowIndex(): number {
    return Math.floor(Date.now() / 1000 / this.windowSize);
  }

  isAllowed(userId: string): boolean {
    const currentWindowIndex = this.getCurrentWindowIndex();
    const record = this.counters.get(userId);

    // 窗口已更換，重置計數
    if (!record || record.windowIndex !== currentWindowIndex) {
      this.counters.set(userId, { count: 1, windowIndex: currentWindowIndex });
      return true;
    }

    if (record.count >= this.limit) {
      return false; // 限流
    }

    record.count++;
    return true;
  }
}

/**
 * Leaky Bucket Rate Limiter
 * - 輸出速率恆定，適合需要平滑流量的場景
 * - 使用佇列模擬，佇列滿時拒絕新請求
 */
class LeakyBucket {
  private capacity: number;     // 佇列最大容量
  private leakRate: number;     // 每秒流出的請求數
  private queue: number[] = []; // 請求進入的時間戳
  private lastLeakTime: number;

  constructor(capacity: number, leakRate: number) {
    this.capacity = capacity;
    this.leakRate = leakRate;
    this.lastLeakTime = Date.now();
  }

  /** 根據經過時間移除已流出的請求 */
  private leak(): void {
    const now = Date.now();
    const elapsedSeconds = (now - this.lastLeakTime) / 1000;
    const leaked = Math.floor(elapsedSeconds * this.leakRate);

    if (leaked > 0) {
      this.queue.splice(0, leaked); // 從前端移除已流出的請求
      this.lastLeakTime = now;
    }
  }

  isAllowed(): boolean {
    this.leak();

    if (this.queue.length >= this.capacity) {
      return false; // 佇列已滿，拒絕
    }

    this.queue.push(Date.now());
    return true;
  }

  getQueueSize(): number {
    this.leak();
    return this.queue.length;
  }
}

/** HTTP Rate Limit 回應結構 */
interface RateLimitResult {
  allowed: boolean;
  limit: number;
  remaining: number;
  resetTime: number;
  retryAfter?: number;
}

/** 通用中介層（適用於 Express / Fastify / NestJS） */
class RateLimiterMiddleware {
  private limiter: SlidingWindowCounter;

  constructor(limit: number = 100, windowSizeSeconds: number = 60) {
    this.limiter = new SlidingWindowCounter(windowSizeSeconds, limit);
    // 每分鐘清理一次過期資料
    setInterval(() => this.limiter.cleanup(), 60_000);
  }

  check(userId: string): RateLimitResult {
    const allowed = this.limiter.isAllowed(userId);
    const remaining = this.limiter.getRemainingRequests(userId);
    const resetTime = this.limiter.getResetTime();

    return {
      allowed,
      limit: 100,
      remaining: allowed ? remaining : 0,
      resetTime,
      retryAfter: allowed
        ? undefined
        : resetTime - Math.floor(Date.now() / 1000),
    };
  }
}

// === 測試所有限流演算法 ===

function runTests(): void {
  console.log("=== Token Bucket Test ===");
  const tbLimiter = new TokenBucketRateLimiter(10, 2); // 容量 10，每秒補充 2

  let allowed = 0;
  let denied = 0;
  for (let i = 0; i < 12; i++) {
    if (tbLimiter.isAllowed("user-1")) {
      allowed++;
    } else {
      denied++;
    }
  }
  console.log(`Allowed: ${allowed}, Denied: ${denied}`);
  // 輸出：Allowed: 10, Denied: 2

  console.log("\n=== Sliding Window Counter Test ===");
  const swLimiter = new SlidingWindowCounter(60, 5); // 每分鐘 5 個

  const results: boolean[] = [];
  for (let i = 0; i < 7; i++) {
    results.push(swLimiter.isAllowed("user-2"));
  }
  console.log(`Results: ${results.map(r => r ? "✓" : "✗").join(" ")}`);
  // 輸出：Results: ✓ ✓ ✓ ✓ ✓ ✗ ✗
  console.log(`Remaining: ${swLimiter.getRemainingRequests("user-2")}`);
  // 輸出：Remaining: 0

  console.log("\n=== Fixed Window Counter Test ===");
  const fwLimiter = new FixedWindowCounter(60, 3); // 每分鐘 3 個

  for (let i = 0; i < 5; i++) {
    const result = fwLimiter.isAllowed("user-3");
    console.log(`Request ${i + 1}: ${result ? "ALLOWED" : "DENIED"}`);
  }
  // 輸出：前 3 個 ALLOWED，後 2 個 DENIED

  console.log("\n=== Leaky Bucket Test ===");
  const lbLimiter = new LeakyBucket(5, 1); // 佇列 5，每秒流出 1 個

  for (let i = 0; i < 7; i++) {
    const result = lbLimiter.isAllowed();
    console.log(`Request ${i + 1}: ${result ? "ALLOWED" : "DENIED"}`);
  }
  // 輸出：前 5 個 ALLOWED，後 2 個 DENIED
}

runTests();

/**
 * High-Performance Token Bucket Rate Limiter (C++)
 *
 * 設計目標：
 * - 使用 CAS（Compare-And-Swap）實現 lock-free 令牌消耗
 * - 多執行緒安全
 * - 最小化鎖競爭
 */

#include <atomic>
#include <chrono>
#include <unordered_map>
#include <mutex>
#include <memory>
#include <iostream>
#include <thread>
#include <vector>

using namespace std;
using namespace std::chrono;

/**
 * 單一用戶的 Token Bucket（執行緒安全）
 * 使用 scaled integer 模擬浮點數原子操作
 */
class AtomicTokenBucket {
private:
    const double capacity_;
    const double refill_rate_;  // 每秒補充令牌數

    atomic<int64_t> tokens_scaled_;   // 實際值 = tokens_scaled_ / SCALE
    atomic<int64_t> last_refill_ns_;  // 上次補充時間（奈秒）

    static constexpr int64_t SCALE = 1000; // 精度 0.001 token

    int64_t now_ns() const {
        return duration_cast<nanoseconds>(
            steady_clock::now().time_since_epoch()
        ).count();
    }

    void refill() {
        int64_t now = now_ns();
        int64_t last = last_refill_ns_.load(memory_order_relaxed);
        int64_t elapsed_ns = now - last;

        if (elapsed_ns <= 0) return;

        double elapsed_sec = static_cast<double>(elapsed_ns) / 1e9;
        int64_t tokens_to_add = static_cast<int64_t>(
            elapsed_sec * refill_rate_ * SCALE
        );

        if (tokens_to_add <= 0) return;

        // CAS 更新上次補充時間
        if (!last_refill_ns_.compare_exchange_strong(
            last, now, memory_order_acq_rel)) {
            return; // 其他執行緒已更新，跳過
        }

        // 補充令牌（不超過容量）
        int64_t cap_scaled = static_cast<int64_t>(capacity_ * SCALE);
        int64_t current = tokens_scaled_.load(memory_order_relaxed);
        int64_t new_tokens = min(cap_scaled, current + tokens_to_add);
        tokens_scaled_.store(new_tokens, memory_order_relaxed);
    }

public:
    AtomicTokenBucket(double capacity, double refill_rate)
        : capacity_(capacity)
        , refill_rate_(refill_rate)
        , tokens_scaled_(static_cast<int64_t>(capacity * SCALE))
        , last_refill_ns_(now_ns()) {}

    /** CAS loop 實現 lock-free 令牌消耗 */
    bool consume(double count = 1.0) {
        refill();

        int64_t cost_scaled = static_cast<int64_t>(count * SCALE);

        int64_t current = tokens_scaled_.load(memory_order_acquire);
        while (true) {
            if (current < cost_scaled) {
                return false; // 令牌不足，拒絕
            }
            int64_t desired = current - cost_scaled;
            if (tokens_scaled_.compare_exchange_weak(
                current, desired,
                memory_order_release,
                memory_order_acquire)) {
                return true; // 成功消耗
            }
            // CAS 失敗：current 已被更新為最新值，自動重試
        }
    }

    double get_tokens() const {
        return static_cast<double>(
            tokens_scaled_.load(memory_order_relaxed)
        ) / SCALE;
    }
};

/** 多用戶 Rate Limiter 管理器 */
class RateLimiterManager {
private:
    unordered_map<string, shared_ptr<AtomicTokenBucket>> buckets_;
    mutable mutex map_mutex_;
    double capacity_;
    double refill_rate_;

public:
    RateLimiterManager(double capacity, double refill_rate)
        : capacity_(capacity), refill_rate_(refill_rate) {}

    bool is_allowed(const string& user_id) {
        shared_ptr<AtomicTokenBucket> bucket;

        {
            lock_guard<mutex> lock(map_mutex_);
            auto it = buckets_.find(user_id);
            if (it == buckets_.end()) {
                auto new_bucket = make_shared<AtomicTokenBucket>(
                    capacity_, refill_rate_
                );
                buckets_[user_id] = new_bucket;
                bucket = new_bucket;
            } else {
                bucket = it->second;
            }
        }

        // bucket 操作不持有 map_mutex_，最大化並行度
        return bucket->consume();
    }

    size_t bucket_count() const {
        lock_guard<mutex> lock(map_mutex_);
        return buckets_.size();
    }
};

/**
 * Sliding Window Counter (C++ 版)
 * - 執行緒安全
 * - 結合 Fixed Window 的低記憶體與 Sliding Window 的平滑性
 */
class SlidingWindowCounter {
private:
    const int window_size_sec_;
    const int limit_;

    struct WindowData {
        atomic<int> current_count{0};
        atomic<int> previous_count{0};
        atomic<int64_t> current_window_index{0};
        mutex window_mutex;
    };

    unordered_map<string, unique_ptr<WindowData>> counters_;
    mutable mutex map_mutex_;

    int64_t get_window_index() const {
        auto now = system_clock::now().time_since_epoch();
        return duration_cast<seconds>(now).count() / window_size_sec_;
    }

public:
    SlidingWindowCounter(int window_size_sec, int limit)
        : window_size_sec_(window_size_sec), limit_(limit) {}

    bool is_allowed(const string& user_id) {
        WindowData* data = nullptr;

        {
            lock_guard<mutex> lock(map_mutex_);
            auto it = counters_.find(user_id);
            if (it == counters_.end()) {
                counters_[user_id] = make_unique<WindowData>();
                data = counters_[user_id].get();
            } else {
                data = it->second.get();
            }
        }

        int64_t current_idx = get_window_index();

        {
            lock_guard<mutex> lock(data->window_mutex);

            // 若窗口已更換，滾動計數器
            int64_t stored_idx = data->current_window_index.load();
            if (stored_idx != current_idx) {
                if (stored_idx + 1 == current_idx) {
                    data->previous_count.store(
                        data->current_count.load()
                    );
                } else {
                    data->previous_count.store(0);
                }
                data->current_count.store(0);
                data->current_window_index.store(current_idx);
            }

            // 計算滑動窗口估算值
            auto now_sec = duration_cast<seconds>(
                system_clock::now().time_since_epoch()
            ).count();
            double progress =
                static_cast<double>(now_sec % window_size_sec_)
                / window_size_sec_;

            double estimated =
                data->previous_count.load() * (1.0 - progress)
                + data->current_count.load();

            if (estimated >= limit_) {
                return false; // 限流
            }

            data->current_count.fetch_add(1, memory_order_relaxed);
            return true;
        }
    }
};

int main() {
    cout << "=== Token Bucket Rate Limiter (C++) ===" << endl;

    RateLimiterManager manager(10.0, 2.0); // 容量 10，每秒補充 2

    // 模擬多執行緒請求
    atomic<int> allowed_count{0};
    atomic<int> denied_count{0};

    auto send_requests = [&](const string& user_id, int count) {
        for (int i = 0; i < count; i++) {
            if (manager.is_allowed(user_id)) {
                allowed_count.fetch_add(1);
            } else {
                denied_count.fetch_add(1);
            }
        }
    };

    // 4 個執行緒，各發送 5 個請求（共 20 個）
    vector<thread> threads;
    for (int i = 0; i < 4; i++) {
        threads.emplace_back(send_requests, "user-1", 5);
    }
    for (auto& t : threads) t.join();

    cout << "Total allowed: " << allowed_count << endl;  // ~10（桶容量）
    cout << "Total denied:  " << denied_count  << endl;  // ~10
    cout << "Buckets in memory: " << manager.bucket_count() << endl;

    cout << "\n=== Sliding Window Counter (C++) ===" << endl;

    SlidingWindowCounter swc(60, 5); // 每分鐘 5 個
    int sw_allowed = 0, sw_denied = 0;

    for (int i = 0; i < 8; i++) {
        if (swc.is_allowed("user-2")) {
            sw_allowed++;
            cout << "Request " << i + 1 << ": ALLOWED" << endl;
        } else {
            sw_denied++;
            cout << "Request " << i + 1 << ": DENIED" << endl;
        }
    }
    cout << "Allowed: " << sw_allowed << ", Denied: " << sw_denied << endl;
    // 輸出：Allowed: 5, Denied: 3

    return 0;
}

-- rate_limit.lua
-- KEYS[1] = 當前窗口 key (e.g., "rl:user123:1716854400")
-- KEYS[2] = 前窗口 key (e.g., "rl:user123:1716854340")
-- ARGV[1] = limit, ARGV[2] = window_size, ARGV[3] = now, ARGV[4] = ttl

local limit = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local ttl = tonumber(ARGV[4])

local current_count = tonumber(redis.call('GET', KEYS[1])) or 0
local previous_count = tonumber(redis.call('GET', KEYS[2])) or 0

-- 計算當前窗口進度
local progress = (now % window_size) / window_size

-- 滑動窗口估算
local estimated = previous_count * (1 - progress) + current_count

if estimated >= limit then
    local reset_at = math.ceil(now / window_size) * window_size
    return {0, limit, 0, reset_at}
end

-- 增加當前窗口計數
local new_count = redis.call('INCR', KEYS[1])

-- 設定 TTL，防止 key 永久留存
if new_count == 1 then
    redis.call('EXPIRE', KEYS[1], ttl)
end

local remaining = limit - math.ceil(
    previous_count * (1 - progress) + new_count
)
local reset_at = math.ceil(now / window_size) * window_size

return {1, limit, math.max(0, remaining), reset_at}

import { createClient } from "redis";

interface RateLimitResponse {
  allowed: boolean;
  limit: number;
  remaining: number;
  resetAt: number;
}

class RedisRateLimiter {
  private client: ReturnType<typeof createClient>;
  private luaScript: string;
  private scriptSha: string | null = null;

  constructor(redisUrl: string = "redis://localhost:6379") {
    this.client = createClient({ url: redisUrl });
    this.luaScript = `
      local limit = tonumber(ARGV[1])
      local window_size = tonumber(ARGV[2])
      local now = tonumber(ARGV[3])
      local ttl = tonumber(ARGV[4])
      local current_count = tonumber(redis.call('GET', KEYS[1])) or 0
      local previous_count = tonumber(redis.call('GET', KEYS[2])) or 0
      local progress = (now % window_size) / window_size
      local estimated = previous_count * (1 - progress) + current_count
      if estimated >= limit then
        local reset_at = math.ceil(now / window_size) * window_size
        return {0, limit, 0, reset_at}
      end
      local new_count = redis.call('INCR', KEYS[1])
      if new_count == 1 then redis.call('EXPIRE', KEYS[1], ttl) end
      local remaining = limit - math.ceil(
        previous_count * (1 - progress) + new_count
      )
      local reset_at = math.ceil(now / window_size) * window_size
      return {1, limit, math.max(0, remaining), reset_at}
    `;
  }

  async connect(): Promise<void> {
    await this.client.connect();
    // 預載 Lua Script，使用 SHA 調用可減少網路傳輸
    this.scriptSha = await this.client.scriptLoad(this.luaScript);
  }

  async isAllowed(
    userId: string,
    limit: number = 100,
    windowSizeSeconds: number = 60
  ): Promise<RateLimitResponse> {
    const now = Math.floor(Date.now() / 1000);
    const currentWindowIndex = Math.floor(now / windowSizeSeconds);
    const previousWindowIndex = currentWindowIndex - 1;

    const currentKey = `rl:${userId}:${currentWindowIndex}`;
    const previousKey = `rl:${userId}:${previousWindowIndex}`;

    const result = await this.client.evalSha(
      this.scriptSha!,
      {
        keys: [currentKey, previousKey],
        arguments: [
          limit.toString(),
          windowSizeSeconds.toString(),
          now.toString(),
          (windowSizeSeconds * 2).toString(), // TTL = 2 個窗口
        ],
      }
    ) as number[];

    return {
      allowed: result[0] === 1,
      limit: result[1],
      remaining: result[2],
      resetAt: result[3],
    };
  }

  /** Express 中介層工廠函數 */
  middleware(limit: number = 100, windowSizeSeconds: number = 60) {
    return async (req: any, res: any, next: any) => {
      const userId = req.headers["x-user-id"] || req.ip;
      const result = await this.isAllowed(userId, limit, windowSizeSeconds);

      // 設定標準 Rate Limit 標頭
      res.set({
        "X-RateLimit-Limit": result.limit,
        "X-RateLimit-Remaining": result.remaining,
        "X-RateLimit-Reset": result.resetAt,
      });

      if (!result.allowed) {
        const retryAfter = result.resetAt - Math.floor(Date.now() / 1000);
        res.set("Retry-After", retryAfter);
        return res.status(429).json({
          error: "Too Many Requests",
          message: `Rate limit exceeded. Try again in ${retryAfter}s.`,
          retryAfter,
        });
      }

      next();
    };
  }
}

async function isAllowedWithFallback(
  userId: string,
  limit: number
): Promise<boolean> {
  try {
    const result = await redisLimiter.isAllowed(userId, limit);
    return result.allowed;
  } catch (error) {
    console.error(
      "Redis Rate Limiter failed, falling back to local",
      error
    );
    // Fail-Local：使用本地 Token Bucket 作為備援
    return localLimiter.isAllowed(userId);
  }
}