Distributed Rate Limits Across Gateways

"""Distributed rate limits across gateway pods: local-only vs global vs local front + global.

Run: python3 distributed_limits.py   (stdlib only, under 1 second)
SharedStore stands in for Redis plus the Lua token bucket from the redis-lua page:
one lock makes check-and-debit atomic, exactly like one EVALSHA call.
"""
import random
import threading
import time


class TokenBucket:
    """Plain in-process token bucket. Not shared: every pod has its own copy."""

    def __init__(self, rate: float, capacity: float) -> None:
        self.rate, self.capacity = rate, capacity
        self.tokens, self.ts = capacity, time.monotonic()  # monotonic: no wall-clock jumps

    def try_take(self, cost: float = 1.0) -> bool:
        now = time.monotonic()
        self.tokens = min(self.capacity, self.tokens + (now - self.ts) * self.rate)
        self.ts = now
        if self.tokens >= cost:
            self.tokens -= cost
            return True
        return False


class SharedStore:
    """Fake Redis: one bucket per key, refill + debit under one lock (the Lua script)."""

    def __init__(self, rate: float, capacity: float) -> None:
        self.rate, self.capacity = rate, capacity
        self.buckets: dict = {}
        self.lock = threading.Lock()
        self.calls = 0          # network round trips we would pay to Redis
        self.down = False       # flip to simulate a Redis outage

    def take(self, key: str, cost: float = 1.0) -> bool:
        if self.down:
            raise TimeoutError("redis timeout")
        with self.lock:         # atomic like EVALSHA: nothing interleaves
            self.calls += 1
            b = self.buckets.setdefault(key, TokenBucket(self.rate, self.capacity))
            return b.try_take(cost)


class GatewayPod:
    """Step 2-6 of Diagram 1. mode is 'local', 'global' or 'hybrid'."""

    def __init__(self, mode: str, store: SharedStore, limit: float, pods: int) -> None:
        self.mode, self.store = mode, store
        if mode == "local":
            # Wrong pick from Diagram 2: every pod enforces the full limit on its own.
            self.local = TokenBucket(rate=limit, capacity=limit)
        else:
            # Hybrid front: a pod's fair share with 50% headroom for uneven load balancing.
            share = limit / pods * 1.5
            self.local = TokenBucket(rate=share, capacity=share)
        # Fallback used only when Redis is down: strict per-pod share, so the total stays near the limit.
        self.fallback = TokenBucket(rate=limit / pods, capacity=limit / pods)
        self.fallbacks = 0

    def handle(self, key: str) -> int:
        if self.mode in ("local", "hybrid") and not self.local.try_take():
            return 429                         # Step 4a: fast reject, no Redis call
        if self.mode == "local":
            return 200
        try:
            ok = self.store.take(key)          # Step 4b: atomic debit on shared Redis
        except TimeoutError:
            self.fallbacks += 1                # Failure path: degrade to local share and alert
            return 200 if self.fallback.try_take() else 429
        return 200 if ok else 429              # Step 5 -> 6a forward, 6b 429 + Retry-After


def run(mode: str, pods: int = 3, limit: float = 100, requests: int = 600, redis_down: bool = False):
    store = SharedStore(rate=limit, capacity=limit)
    store.down = redis_down
    fleet = [GatewayPod(mode, store, limit, pods) for _ in range(pods)]
    rng = random.Random(7)
    codes = [rng.choice(fleet).handle("tenant:42") for _ in range(requests)]  # LB spreads one key
    return codes.count(200), store.calls, sum(p.fallbacks for p in fleet)


if __name__ == "__main__":
    print("limit 100 per s for key tenant:42, 3 pods, 600 requests in one burst")
    for mode in ("local", "global", "hybrid"):
        ok, calls, _ = run(mode)
        print(f"{mode:7s} admitted={ok:3d} redis_calls={calls}")
    ok, calls, fb = run("hybrid", redis_down=True)
    print(f"hybrid with Redis down: admitted={ok} redis_calls={calls} fallbacks={fb}")
    ok6, _, _ = run("local", pods=6)
    print(f"local with 6 pods: admitted={ok6}  (limit multiplies with pod count)")

/**
 * Distributed rate limits: local token bucket in front of a shared (Redis-like) store.
 * Run: npx tsx distributed_limits.ts   (no deps; the store is simulated with 2 ms latency)
 * The fake store does refill + debit in one step, standing in for one EVALSHA Lua call.
 */

class TokenBucket {
  private tokens: number;
  private ts = performance.now(); // monotonic clock, safe against wall-clock jumps
  constructor(private rate: number, private capacity: number) {
    this.tokens = capacity;
  }
  tryTake(cost = 1): boolean {
    const now = performance.now();
    this.tokens = Math.min(this.capacity, this.tokens + ((now - this.ts) / 1000) * this.rate);
    this.ts = now;
    if (this.tokens >= cost) {
      this.tokens -= cost;
      return true;
    }
    return false;
  }
}

const sleep = (ms: number) => new Promise<void>((r) => setTimeout(r, ms));

class FakeRedis {
  calls = 0;
  down = false; // when true, calls hang until the client timeout fires
  private buckets = new Map<string, TokenBucket>();
  constructor(private rate: number, private capacity: number) {}
  async take(key: string): Promise<boolean> {
    if (this.down) return new Promise<boolean>(() => {}); // never resolves, like a dead socket
    await sleep(2); // network hop
    this.calls++;
    let b = this.buckets.get(key);
    if (!b) this.buckets.set(key, (b = new TokenBucket(this.rate, this.capacity)));
    return b.tryTake(); // single-threaded event loop: this block is atomic, like Lua
  }
}

function withTimeout<T>(p: Promise<T>, ms: number): Promise<T> {
  return Promise.race([p, sleep(ms).then(() => Promise.reject(new Error("redis timeout")))]);
}

class GatewayPod {
  private front: TokenBucket;    // Step 2: per-pod share with headroom
  private fallback: TokenBucket; // strict share, used only while Redis is down
  fallbacks = 0;
  constructor(private redis: FakeRedis, limit: number, pods: number) {
    this.front = new TokenBucket((limit / pods) * 1.5, (limit / pods) * 1.5);
    this.fallback = new TokenBucket(limit / pods, limit / pods);
  }
  async handle(key: string): Promise<number> {
    if (!this.front.tryTake()) return 429; // Step 4a: fast 429, no Redis call
    try {
      const ok = await withTimeout(this.redis.take(key), 20); // Step 4b: global debit
      return ok ? 200 : 429; // Step 5 -> 6a forward or 6b 429 with Retry-After
    } catch {
      this.fallbacks++; // Failure path: local limit only, page someone
      return this.fallback.tryTake() ? 200 : 429;
    }
  }
}

async function run(redisDown: boolean) {
  const limit = 100, pods = 3;
  const redis = new FakeRedis(limit, limit);
  redis.down = redisDown;
  const fleet = Array.from({ length: pods }, () => new GatewayPod(redis, limit, pods));
  const codes = await Promise.all(
    Array.from({ length: 600 }, (_, i) => fleet[i % pods]!.handle("tenant:42")),
  );
  const ok = codes.filter((c) => c === 200).length;
  const fb = fleet.reduce((s, p) => s + p.fallbacks, 0);
  console.log(`redisDown=${redisDown} admitted=${ok} redisCalls=${redis.calls} fallbacks=${fb}`);
}

await run(false);
await run(true);