"""Reference solution — Rate Limiter: Bucket to Distributed. DO NOT READ BEFORE YOU HAVE RUN THE PROBLEM UNDER THE CLOCK. Teaching text: ../../../WARMUP.md § Chapter 4. The four algorithms and what each actually costs: Fixed window O(1), and admits 2x the limit across a boundary. Sliding window log exact, O(limit) memory PER KEY. The correctness baseline. Sliding window ctr O(1), no boundary burst, ~1% error on real traffic. Token bucket O(1), and it can express "average 10/s but 100 at once is fine" -- which no window can. The design decision that matters most in all four: the clock is INJECTED. A limiter you cannot test deterministically is a limiter you cannot ship, because the only alternative is tests that sleep, and tests that sleep get deleted. """ from __future__ import annotations import threading import time from collections import deque class StoreUnavailable(Exception): """The shared counter store could not be reached.""" # --------------------------------------------------------------------------- # Gate 1 # --------------------------------------------------------------------------- class TokenBucket: """Sustained `rate` per second, burst up to `capacity`. Lazy refill.""" __slots__ = ("capacity", "rate", "_tokens", "_last", "_clock") def __init__(self, capacity, rate, clock=time.monotonic): if capacity <= 0: raise ValueError("capacity must be positive") if rate <= 0: raise ValueError("rate must be positive") self.capacity = float(capacity) self.rate = float(rate) self._tokens = float(capacity) self._clock = clock self._last = clock() def _refill(self): now = self._clock() elapsed = now - self._last if elapsed > 0: # No background thread: derive the refill from elapsed time. This # is what makes the bucket O(1) and testable with a fake clock. self._tokens = min(self.capacity, self._tokens + elapsed * self.rate) self._last = now def allow(self, cost=1): self._refill() if self._tokens >= cost: self._tokens -= cost return True return False def retry_after(self, cost=1): """Seconds until `cost` tokens exist. Send this in the 429.""" self._refill() if self._tokens >= cost: return 0.0 return (cost - self._tokens) / self.rate @property def tokens(self): self._refill() return self._tokens # --------------------------------------------------------------------------- # Gate 2 # --------------------------------------------------------------------------- class SlidingWindowLog: """Exactly correct. O(limit) memory per key — the correctness baseline.""" __slots__ = ("limit", "window", "_log", "_clock") def __init__(self, limit, window, clock=time.monotonic): if limit <= 0 or window <= 0: raise ValueError("limit and window must be positive") self.limit, self.window = limit, float(window) self._log = deque() self._clock = clock def allow(self): now = self._clock() cutoff = now - self.window while self._log and self._log[0] <= cutoff: self._log.popleft() if len(self._log) < self.limit: self._log.append(now) return True return False class SlidingWindowCounter: """O(1) memory, no boundary burst, ~1% error on real traffic. Two counters -- current and previous fixed window -- and an interpolation weighted by how far into the current window we are. The lie you must disclose: it assumes the previous window's requests were spread uniformly. """ __slots__ = ("limit", "window", "_curr_start", "_curr", "_prev", "_clock") def __init__(self, limit, window, clock=time.monotonic): if limit <= 0 or window <= 0: raise ValueError("limit and window must be positive") self.limit, self.window = limit, float(window) self._clock = clock self._curr_start = self._floor(clock()) self._curr = self._prev = 0 def _floor(self, now): return now - (now % self.window) def _roll(self, now): start = self._floor(now) if start == self._curr_start: return if abs(start - self._curr_start - self.window) < 1e-9: self._prev, self._curr = self._curr, 0 # slid by exactly one else: self._prev, self._curr = 0, 0 # gap: both are stale self._curr_start = start def allow(self): now = self._clock() self._roll(now) weight = 1.0 - (now - self._curr_start) / self.window if self._prev * weight + self._curr < self.limit: self._curr += 1 return True return False # --------------------------------------------------------------------------- # Gate 3 # --------------------------------------------------------------------------- class ShardedLimiter: """Per-key limiters with sharded locks. One global lock is correct and is the bottleneck, because EVERY call mutates state -- there is no read-only fast path to optimise. Sharding by hash(key) removes the contention; the cost is that shard count is now a tuning parameter that should come from measured contention rather than from a round number. """ def __init__(self, factory, shards=16): if shards <= 0: raise ValueError("shards must be positive") self._factory = factory self._shards = [({}, threading.Lock()) for _ in range(shards)] def _shard(self, key): return self._shards[hash(key) % len(self._shards)] def allow(self, key, cost=1): limiters, lock = self._shard(key) with lock: limiter = limiters.get(key) if limiter is None: limiter = limiters[key] = self._factory() return limiter.allow(cost) def __len__(self): return sum(len(m) for m, _ in self._shards) # --------------------------------------------------------------------------- # Gate 4 # --------------------------------------------------------------------------- class DistributedLimiter: """One limit shared across processes, via an atomic counter store. Four problems, all of which must be named: 1. Atomicity -- read-then-write over the network races. incr() must be one atomic operation on the store. 2. Round trips-- a network hop per request can cost more than the work being protected. `lease` claims a batch and spends it locally: precision traded for latency. 3. Clock skew -- the window is derived from the caller's clock here, but in production you would use the store's clock (Redis TIME) so all callers agree. 4. Availability- the store WILL be down. fail_open admits (right for overload protection); fail_closed denies (right for billing and abuse). Neither is universally correct, and the sophisticated answer is fail-open with a degraded local limit of global/process_count. """ def __init__(self, store, limit, window, clock=time.monotonic, fail_open=True, lease=1): if limit <= 0 or window <= 0: raise ValueError("limit and window must be positive") if lease < 1: raise ValueError("lease must be >= 1") self._store = store self.limit, self.window = limit, float(window) self._clock = clock self.fail_open = fail_open self.lease = lease self._local = {} # key -> [window_index, permits_remaining] self._lock = threading.Lock() def _window_index(self, now): return int(now // self.window) def allow(self, key): now = self._clock() index = self._window_index(now) with self._lock: held = self._local.get(key) if held is not None and held[0] == index and held[1] > 0: held[1] -= 1 # spend a locally-held permit return True bucket = f"{key}:{index}" try: # Claim `lease` permits in ONE atomic round trip. count = self._store.incr(bucket, self.window * 2, self.lease) except StoreUnavailable: return self.fail_open # `count` is the total claimed in this window, including ours. Work out # how many of the ones we just claimed are actually within budget. over = count - self.limit granted = self.lease if over <= 0 else max(0, self.lease - over) with self._lock: if granted <= 0: self._local[key] = [index, 0] return False self._local[key] = [index, granted - 1] return True