Design a Rate Limiter.

# rate_limiter.py, part 1: time, the decision, and the one interface every algorithm implements
from __future__ import annotations

import threading
import time
from abc import ABC, abstractmethod
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass


class Clock(ABC):
    @abstractmethod
    def now(self) -> float: ...  # seconds; must never go backwards


class MonotonicClock(Clock):
    def now(self) -> float:
        return time.monotonic()


class FakeClock(Clock):
    def __init__(self, start: float = 0.0) -> None:
        self.t = start

    def now(self) -> float:
        return self.t

    def advance(self, seconds: float) -> None:
        self.t += seconds


@dataclass(frozen=True)
class Decision:
    allowed: bool
    remaining: int
    retry_after: float  # seconds until a retry can succeed; 0 when allowed


class RateLimiter(ABC):
    """One limiter guards one key (a user, an API key, an IP)."""

    @abstractmethod
    def try_acquire(self, cost: int = 1) -> Decision: ...

# part 2: three algorithms behind the same interface
class TokenBucket(RateLimiter):
    def __init__(self, capacity: int, refill_per_sec: float, clock: Clock) -> None:
        self.capacity, self.rate, self.clock = capacity, refill_per_sec, clock
        self.tokens = float(capacity)
        self.last = clock.now()
        self._lock = threading.Lock()

    def try_acquire(self, cost: int = 1) -> Decision:
        if cost > self.capacity:
            raise ValueError(f"cost {cost} can never fit in a bucket of {self.capacity}")
        with self._lock:
            now = self.clock.now()
            self.tokens = min(self.capacity, self.tokens + (now - self.last) * self.rate)  # lazy refill
            self.last = now
            if self.tokens >= cost:
                self.tokens -= cost
                return Decision(True, int(self.tokens), 0.0)
            return Decision(False, int(self.tokens), (cost - self.tokens) / self.rate)


class FixedWindow(RateLimiter):
    def __init__(self, limit: int, window: float, clock: Clock) -> None:
        self.limit, self.window, self.clock = limit, window, clock
        self.start, self.count = 0.0, 0
        self._lock = threading.Lock()

    def try_acquire(self, cost: int = 1) -> Decision:
        with self._lock:
            now = self.clock.now()
            if now - self.start >= self.window:
                self.start, self.count = now - (now % self.window), 0
            if self.count + cost <= self.limit:
                self.count += cost
                return Decision(True, self.limit - self.count, 0.0)
            return Decision(False, 0, self.start + self.window - now)


class SlidingWindowLog(RateLimiter):
    def __init__(self, limit: int, window: float, clock: Clock) -> None:
        self.limit, self.window, self.clock = limit, window, clock
        self.log: deque[float] = deque()
        self._lock = threading.Lock()

    def try_acquire(self, cost: int = 1) -> Decision:
        with self._lock:
            now = self.clock.now()
            while self.log and self.log[0] <= now - self.window:
                self.log.popleft()
            if len(self.log) + cost <= self.limit:
                self.log.extend([now] * cost)
                return Decision(True, self.limit - len(self.log), 0.0)
            return Decision(False, 0, self.log[0] + self.window - now)

# part 3: rules, a factory that builds limiters from them, and the per-key registry
@dataclass(frozen=True)
class Rule:
    algorithm: str  # "token_bucket" | "fixed_window" | "sliding_log"
    limit: int
    window: float  # seconds


LimiterFactory = Callable[[Rule, Clock], RateLimiter]

FACTORIES: dict[str, LimiterFactory] = {
    "token_bucket": lambda r, c: TokenBucket(r.limit, r.limit / r.window, c),
    "fixed_window": lambda r, c: FixedWindow(r.limit, r.window, c),
    "sliding_log": lambda r, c: SlidingWindowLog(r.limit, r.window, c),
}


class RateLimiterRegistry:
    """Finds the rule for a (client tier, endpoint) and keeps one limiter per client key."""

    def __init__(self, rules: dict[tuple[str, str], Rule], default: Rule, clock: Clock) -> None:
        self.rules, self.default, self.clock = rules, default, clock
        self._limiters: dict[tuple[str, str], RateLimiter] = {}
        self._lock = threading.Lock()  # only guards creation; each limiter has its own lock

    def _limiter(self, client: str, tier: str, endpoint: str) -> RateLimiter:
        key = (client, endpoint)
        lim = self._limiters.get(key)
        if lim is None:
            with self._lock:
                lim = self._limiters.get(key)  # check again: another thread may have won
                if lim is None:
                    rule = self.rules.get((tier, endpoint), self.default)
                    lim = FACTORIES[rule.algorithm](rule, self.clock)
                    self._limiters[key] = lim
        return lim

    def check(self, client: str, tier: str, endpoint: str, cost: int = 1) -> Decision:
        return self._limiter(client, tier, endpoint).try_acquire(cost)

# part 4: the behaviour, checked with a fake clock
if __name__ == "__main__":
    clock = FakeClock(start=1000.0)
    rules = {
        ("free", "/otp"): Rule("fixed_window", 3, 60),
        ("free", "/search"): Rule("token_bucket", 5, 1),
        ("pro", "/search"): Rule("token_bucket", 50, 1),
    }
    reg = RateLimiterRegistry(rules, Rule("sliding_log", 10, 60), clock)

    burst = [reg.check("u1", "free", "/search").allowed for _ in range(7)]
    print("burst of 7:", burst)  # 5 True, then False False
    clock.advance(0.5)  # 0.5 s x 5 tokens/s = 2.5 tokens back
    print("after 0.5s:", [reg.check("u1", "free", "/search").allowed for _ in range(3)])  # True True False

    otp = [reg.check("u2", "free", "/otp") for _ in range(4)]
    print("otp:", [d.allowed for d in otp], "retry in", round(otp[-1].retry_after, 1), "s")  # 3 True, 1 False, 19.5 s

    # fixed window's flaw: 3 at the end of one window + 3 at the start of the next = 6 in 2 seconds
    fw = FixedWindow(3, 60, clock)
    clock.t = 1199.0
    a = sum(fw.try_acquire().allowed for _ in range(3))
    clock.t = 1201.0
    b = sum(fw.try_acquire().allowed for _ in range(3))
    sl = SlidingWindowLog(3, 60, clock)
    clock.t = 1199.0
    c = sum(sl.try_acquire().allowed for _ in range(3))
    clock.t = 1201.0
    d = sum(sl.try_acquire().allowed for _ in range(3))
    print("fixed window let through", a + b, "in 2 s; sliding log", c + d)  # 6 vs 3

    # 20 threads, 100 calls each, against one bucket of 50 with no time passing: exactly 50 allowed
    shared = TokenBucket(50, 1.0, FakeClock())
    allowed = []

    def hammer() -> None:
        allowed.append(sum(shared.try_acquire().allowed for _ in range(100)))

    ts = [threading.Thread(target=hammer) for _ in range(20)]
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    print("allowed under contention:", sum(allowed))  # 50

burst of 7: [True, True, True, True, True, False, False]
after 0.5s: [True, True, False]
otp: [True, True, True, False] retry in 19.5 s
fixed window let through 6 in 2 s; sliding log 3
allowed under contention: 50