feat(core): add resilient fixed-window throttling
This commit is contained in:
349
src/govoplan_core/core/throttling.py
Normal file
349
src/govoplan_core/core/throttling.py
Normal file
@@ -0,0 +1,349 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Protocol
|
||||
|
||||
from redis import Redis
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_NAMESPACE_PATTERN = re.compile(r"^[a-z0-9][a-z0-9._-]{0,99}$")
|
||||
_REDIS_INCREMENT_SCRIPT = """
|
||||
local count = redis.call('INCR', KEYS[1])
|
||||
if count == 1 then
|
||||
redis.call('EXPIRE', KEYS[1], ARGV[1])
|
||||
end
|
||||
local ttl = redis.call('TTL', KEYS[1])
|
||||
return {count, ttl}
|
||||
"""
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class FixedWindowBucket:
|
||||
count: int = 0
|
||||
retry_after_seconds: int = 0
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ThrottleDimension:
|
||||
"""One independently limited dimension without putting its value in keys."""
|
||||
|
||||
namespace: str
|
||||
subject: str
|
||||
limit: int
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ThrottleDecision:
|
||||
allowed: bool
|
||||
retry_after_seconds: int = 0
|
||||
|
||||
|
||||
class FixedWindowStore(Protocol):
|
||||
def read(self, key: str) -> FixedWindowBucket: ...
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket: ...
|
||||
|
||||
def delete(self, key: str) -> None: ...
|
||||
|
||||
|
||||
class InMemoryFixedWindowStore:
|
||||
"""Bounded process-local store for development and distributed-store loss."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
max_entries: int = 10_000,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._entries: dict[str, tuple[int, float]] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._max_entries = max(2, max_entries)
|
||||
self._clock = clock
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
entry = self._active_entry(key, now=now)
|
||||
if entry is None:
|
||||
return FixedWindowBucket()
|
||||
count, expires_at = entry
|
||||
return FixedWindowBucket(
|
||||
count=count,
|
||||
retry_after_seconds=max(1, int(expires_at - now + 0.999)),
|
||||
)
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
now = self._clock()
|
||||
with self._lock:
|
||||
entry = self._active_entry(key, now=now)
|
||||
if entry is None:
|
||||
self._make_room(now=now, incoming_key=key)
|
||||
count = 1
|
||||
expires_at = now + window_seconds
|
||||
else:
|
||||
count = entry[0] + 1
|
||||
expires_at = entry[1]
|
||||
self._entries[key] = (count, expires_at)
|
||||
return FixedWindowBucket(
|
||||
count=count,
|
||||
retry_after_seconds=max(1, int(expires_at - now + 0.999)),
|
||||
)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
with self._lock:
|
||||
self._entries.pop(key, None)
|
||||
|
||||
def _active_entry(self, key: str, *, now: float) -> tuple[int, float] | None:
|
||||
entry = self._entries.get(key)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry[1] <= now:
|
||||
self._entries.pop(key, None)
|
||||
return None
|
||||
return entry
|
||||
|
||||
def _make_room(self, *, now: float, incoming_key: str) -> None:
|
||||
if incoming_key in self._entries or len(self._entries) < self._max_entries:
|
||||
return
|
||||
for key, (_count, expires_at) in tuple(self._entries.items()):
|
||||
if expires_at <= now:
|
||||
self._entries.pop(key, None)
|
||||
while len(self._entries) >= self._max_entries:
|
||||
self._entries.pop(next(iter(self._entries)))
|
||||
|
||||
|
||||
class RedisFixedWindowStore:
|
||||
"""Atomic Redis fixed-window counters shared across API workers."""
|
||||
|
||||
def __init__(self, redis_url: str) -> None:
|
||||
self._client = Redis.from_url(
|
||||
redis_url,
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=0.25,
|
||||
socket_timeout=0.25,
|
||||
health_check_interval=30,
|
||||
)
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
pipeline = self._client.pipeline(transaction=False)
|
||||
pipeline.get(key)
|
||||
pipeline.ttl(key)
|
||||
raw_count, raw_ttl = pipeline.execute()
|
||||
return FixedWindowBucket(
|
||||
count=int(raw_count or 0),
|
||||
retry_after_seconds=max(0, int(raw_ttl or 0)),
|
||||
)
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
result = self._client.eval(
|
||||
_REDIS_INCREMENT_SCRIPT,
|
||||
1,
|
||||
key,
|
||||
window_seconds,
|
||||
)
|
||||
if not isinstance(result, (list, tuple)) or len(result) != 2:
|
||||
raise RedisError("Unexpected fixed-window throttle response from Redis")
|
||||
return FixedWindowBucket(
|
||||
count=int(result[0]),
|
||||
retry_after_seconds=max(1, int(result[1])),
|
||||
)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self._client.delete(key)
|
||||
|
||||
|
||||
class ResilientFixedWindowStore:
|
||||
"""Mirror locally, prefer Redis, and fall back safely during an outage."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
primary: FixedWindowStore | None,
|
||||
fallback: FixedWindowStore,
|
||||
*,
|
||||
retry_seconds: int = 30,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._primary = primary
|
||||
self._fallback = fallback
|
||||
self._retry_seconds = max(1, retry_seconds)
|
||||
self._clock = clock
|
||||
self._primary_unavailable_until = 0.0
|
||||
self._state_lock = threading.Lock()
|
||||
|
||||
def read(self, key: str) -> FixedWindowBucket:
|
||||
fallback_result = self._fallback.read(key)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return fallback_result
|
||||
try:
|
||||
return _stricter_bucket(primary.read(key), fallback_result)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
return fallback_result
|
||||
|
||||
def increment(self, key: str, *, window_seconds: int) -> FixedWindowBucket:
|
||||
fallback_result = self._fallback.increment(
|
||||
key,
|
||||
window_seconds=window_seconds,
|
||||
)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return fallback_result
|
||||
try:
|
||||
primary_result = primary.increment(
|
||||
key,
|
||||
window_seconds=window_seconds,
|
||||
)
|
||||
return _stricter_bucket(primary_result, fallback_result)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
return fallback_result
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
self._fallback.delete(key)
|
||||
primary = self._available_primary()
|
||||
if primary is None:
|
||||
return
|
||||
try:
|
||||
primary.delete(key)
|
||||
except (RedisError, OSError, TimeoutError, ConnectionError) as exc:
|
||||
self._mark_primary_unavailable(exc)
|
||||
|
||||
def _available_primary(self) -> FixedWindowStore | None:
|
||||
if self._primary is None:
|
||||
return None
|
||||
with self._state_lock:
|
||||
if self._clock() < self._primary_unavailable_until:
|
||||
return None
|
||||
return self._primary
|
||||
|
||||
def _mark_primary_unavailable(self, exc: Exception) -> None:
|
||||
should_log = False
|
||||
with self._state_lock:
|
||||
now = self._clock()
|
||||
if now >= self._primary_unavailable_until:
|
||||
should_log = True
|
||||
self._primary_unavailable_until = now + self._retry_seconds
|
||||
if should_log:
|
||||
logger.warning(
|
||||
"Distributed throttling is unavailable; using the process-local fallback (%s)",
|
||||
type(exc).__name__,
|
||||
)
|
||||
|
||||
|
||||
def _stricter_bucket(
|
||||
first: FixedWindowBucket,
|
||||
second: FixedWindowBucket,
|
||||
) -> FixedWindowBucket:
|
||||
return FixedWindowBucket(
|
||||
count=max(first.count, second.count),
|
||||
retry_after_seconds=max(
|
||||
first.retry_after_seconds,
|
||||
second.retry_after_seconds,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
class FixedWindowThrottle:
|
||||
def __init__(
|
||||
self,
|
||||
store: FixedWindowStore,
|
||||
*,
|
||||
window_seconds: int,
|
||||
key_prefix: str = "govoplan:throttle:v1",
|
||||
) -> None:
|
||||
self._store = store
|
||||
self._window_seconds = max(1, window_seconds)
|
||||
self._key_prefix = key_prefix.rstrip(":")
|
||||
|
||||
def check(self, dimensions: Iterable[ThrottleDimension]) -> ThrottleDecision:
|
||||
return self._decision(tuple(dimensions), increment=False)
|
||||
|
||||
def record(self, dimensions: Iterable[ThrottleDimension]) -> ThrottleDecision:
|
||||
return self._decision(tuple(dimensions), increment=True)
|
||||
|
||||
def reset(self, dimensions: Iterable[ThrottleDimension]) -> None:
|
||||
for dimension in tuple(dimensions):
|
||||
self._store.delete(self._key(dimension))
|
||||
|
||||
def _decision(
|
||||
self,
|
||||
dimensions: tuple[ThrottleDimension, ...],
|
||||
*,
|
||||
increment: bool,
|
||||
) -> ThrottleDecision:
|
||||
if not dimensions:
|
||||
raise ValueError("At least one throttle dimension is required")
|
||||
blocked_retry_after = 0
|
||||
for dimension in dimensions:
|
||||
if dimension.limit < 1:
|
||||
raise ValueError("Throttle limits must be positive")
|
||||
key = self._key(dimension)
|
||||
state = (
|
||||
self._store.increment(key, window_seconds=self._window_seconds)
|
||||
if increment
|
||||
else self._store.read(key)
|
||||
)
|
||||
if state.count >= dimension.limit:
|
||||
blocked_retry_after = max(
|
||||
blocked_retry_after,
|
||||
max(1, state.retry_after_seconds),
|
||||
)
|
||||
return ThrottleDecision(
|
||||
allowed=blocked_retry_after == 0,
|
||||
retry_after_seconds=blocked_retry_after,
|
||||
)
|
||||
|
||||
def _key(self, dimension: ThrottleDimension) -> str:
|
||||
namespace = dimension.namespace.strip().casefold()
|
||||
if not _NAMESPACE_PATTERN.fullmatch(namespace):
|
||||
raise ValueError("Throttle namespaces must use lowercase letters, digits, '.', '_' or '-'")
|
||||
subject_digest = hashlib.sha256(dimension.subject.encode("utf-8")).hexdigest()
|
||||
return f"{self._key_prefix}:{namespace}:{subject_digest}"
|
||||
|
||||
|
||||
def build_fixed_window_throttle(
|
||||
*,
|
||||
redis_url: str | None,
|
||||
window_seconds: int,
|
||||
redis_retry_seconds: int = 30,
|
||||
max_local_entries: int = 10_000,
|
||||
key_prefix: str = "govoplan:throttle:v1",
|
||||
) -> FixedWindowThrottle:
|
||||
primary = (
|
||||
RedisFixedWindowStore(redis_url)
|
||||
if redis_url is not None and redis_url.strip()
|
||||
else None
|
||||
)
|
||||
store = ResilientFixedWindowStore(
|
||||
primary,
|
||||
InMemoryFixedWindowStore(max_entries=max_local_entries),
|
||||
retry_seconds=redis_retry_seconds,
|
||||
)
|
||||
return FixedWindowThrottle(
|
||||
store,
|
||||
window_seconds=window_seconds,
|
||||
key_prefix=key_prefix,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"FixedWindowBucket",
|
||||
"FixedWindowStore",
|
||||
"FixedWindowThrottle",
|
||||
"InMemoryFixedWindowStore",
|
||||
"RedisFixedWindowStore",
|
||||
"ResilientFixedWindowStore",
|
||||
"ThrottleDecision",
|
||||
"ThrottleDimension",
|
||||
"build_fixed_window_throttle",
|
||||
]
|
||||
Reference in New Issue
Block a user