LLD Case Study: Extensible Rate Limiter

from __future__ import annotations
import threading
import time
from abc import ABC, abstractmethod
from collections import deque
from dataclasses import dataclass
from typing import Optional


# ── Value objects ──────────────────────────────────────────────────────────

@dataclass(frozen=True)
class RatePolicy:
    limit: int       # max requests
    window_ms: int   # window duration in milliseconds


@dataclass(frozen=True)
class Decision:
    allowed: bool
    remaining: int           # tokens/requests remaining (0 if denied)
    retry_after_ms: int      # ms until a request will be allowed (0 if allowed)


# ── Clock abstraction — critical for deterministic testing ─────────────────

class Clock(ABC):
    @abstractmethod
    def now_ms(self) -> int: ...

class SystemClock(Clock):
    def now_ms(self) -> int:
        return int(time.time() * 1000)

class FakeClock(Clock):
    def __init__(self, start_ms: int = 0):
        self._now = start_ms
    def now_ms(self) -> int:
        return self._now
    def advance(self, delta_ms: int) -> None:
        self._now += delta_ms


# ── Store abstraction ──────────────────────────────────────────────────────

class Store(ABC):
    @abstractmethod
    def get(self, key: str) -> Optional[str]: ...
    @abstractmethod
    def set(self, key: str, value: str, ttl_ms: int) -> None: ...
    @abstractmethod
    def delete(self, key: str) -> None: ...


class InMemoryStore(Store):
    """Thread-safe in-memory store for local/test use."""

    def __init__(self):
        self._data: dict[str, tuple[str, Optional[int]]] = {}  # key → (value, expiry_ms)
        self._lock = threading.Lock()

    def get(self, key: str) -> Optional[str]:
        with self._lock:
            entry = self._data.get(key)
            if entry is None:
                return None
            value, expiry_ms = entry
            if expiry_ms is not None and int(time.time() * 1000) > expiry_ms:
                del self._data[key]
                return None
            return value

    def set(self, key: str, value: str, ttl_ms: int) -> None:
        with self._lock:
            expiry = int(time.time() * 1000) + ttl_ms if ttl_ms > 0 else None
            self._data[key] = (value, expiry)

    def delete(self, key: str) -> None:
        with self._lock:
            self._data.pop(key, None)


# ── Algorithm interface ────────────────────────────────────────────────────

class Algorithm(ABC):
    @abstractmethod
    def allow(self, key: str, policy: RatePolicy, clock: Clock) -> Decision: ...
    @abstractmethod
    def reset(self, key: str) -> None: ...


# ── Token Bucket ───────────────────────────────────────────────────────────

class TokenBucketAlgorithm(Algorithm):
    """
    Bucket capacity = policy.limit (burst size).
    Refill rate = policy.limit tokens per policy.window_ms milliseconds.
    Each request consumes 1 token.
    State per key: (tokens: float, last_refill_ms: int) stored as "tokens:timestamp".
    """

    def __init__(self, store: Store):
        self._store = store
        self._locks: dict[str, threading.Lock] = {}
        self._meta_lock = threading.Lock()

    def _key_lock(self, key: str) -> threading.Lock:
        with self._meta_lock:
            if key not in self._locks:
                self._locks[key] = threading.Lock()
            return self._locks[key]

    def allow(self, key: str, policy: RatePolicy, clock: Clock) -> Decision:
        bucket_key = f"tb:{key}"
        with self._key_lock(bucket_key):
            now_ms = clock.now_ms()
            raw = self._store.get(bucket_key)

            if raw is None:
                tokens = float(policy.limit)
                last_refill = now_ms
            else:
                parts = raw.split(":")
                tokens, last_refill = float(parts[0]), int(parts[1])

            # Refill based on elapsed time
            elapsed_ms = now_ms - last_refill
            refill_rate = policy.limit / policy.window_ms  # tokens per ms
            tokens = min(float(policy.limit), tokens + elapsed_ms * refill_rate)
            last_refill = now_ms

            if tokens >= 1.0:
                tokens -= 1.0
                self._store.set(bucket_key, f"{tokens}:{last_refill}", policy.window_ms * 2)
                return Decision(allowed=True, remaining=int(tokens), retry_after_ms=0)
            else:
                self._store.set(bucket_key, f"{tokens}:{last_refill}", policy.window_ms * 2)
                # Time until 1 token is available
                wait_ms = int((1.0 - tokens) / refill_rate)
                return Decision(allowed=False, remaining=0, retry_after_ms=wait_ms)

    def reset(self, key: str) -> None:
        self._store.delete(f"tb:{key}")


# ── Sliding Window Log ─────────────────────────────────────────────────────

class SlidingWindowLogAlgorithm(Algorithm):
    """
    Exact sliding window: maintains a deque of request timestamps per key.
    Memory cost: O(limit) per key (at most `limit` timestamps in the window).
    Correct across window boundaries — no burst at window edges.
    """

    def __init__(self):
        self._logs: dict[str, deque[int]] = {}
        self._locks: dict[str, threading.Lock] = {}
        self._meta_lock = threading.Lock()

    def _key_lock(self, key: str) -> threading.Lock:
        with self._meta_lock:
            if key not in self._locks:
                self._locks[key] = threading.Lock()
            return self._locks[key]

    def allow(self, key: str, policy: RatePolicy, clock: Clock) -> Decision:
        with self._key_lock(key):
            now_ms = clock.now_ms()
            window_start = now_ms - policy.window_ms

            if key not in self._logs:
                self._logs[key] = deque()

            log = self._logs[key]

            # Evict timestamps outside the sliding window
            while log and log[0] <= window_start:
                log.popleft()

            if len(log) < policy.limit:
                log.append(now_ms)
                remaining = policy.limit - len(log)
                return Decision(allowed=True, remaining=remaining, retry_after_ms=0)
            else:
                # Earliest request will exit window at: log[0] + window_ms
                retry_at = log[0] + policy.window_ms - now_ms
                return Decision(allowed=False, remaining=0, retry_after_ms=max(0, retry_at))

    def reset(self, key: str) -> None:
        with self._meta_lock:
            self._logs.pop(key, None)
            self._locks.pop(key, None)


# ── PolicyResolver ─────────────────────────────────────────────────────────

class PolicyResolver:
    """Maps a key to its rate policy. Supports per-tenant overrides."""

    def __init__(self, default: RatePolicy, overrides: Optional[dict[str, RatePolicy]] = None):
        self._default = default
        self._overrides = overrides or {}

    def resolve(self, key: str) -> RatePolicy:
        return self._overrides.get(key, self._default)


# ── RateLimiter — thin orchestrator ────────────────────────────────────────

class RateLimiter:
    def __init__(
        self,
        algorithm: Algorithm,
        clock: Clock,
        policy_resolver: PolicyResolver,
        fail_open: bool = True,  # fail-open: allow requests when limiter errors
    ):
        self._algorithm = algorithm
        self._clock = clock
        self._policy_resolver = policy_resolver
        self._fail_open = fail_open

    def allow(self, key: str) -> Decision:
        try:
            policy = self._policy_resolver.resolve(key)
            return self._algorithm.allow(key, policy, self._clock)
        except Exception:
            # Storage/clock failure: apply failure policy
            if self._fail_open:
                return Decision(allowed=True, remaining=-1, retry_after_ms=0)
            else:
                return Decision(allowed=False, remaining=0, retry_after_ms=1000)

    def reset(self, key: str) -> None:
        self._algorithm.reset(key)