Initial commit

yksa_kit/host_rate_limit.py

@@ -0,0 +1,81 @@
+"""Redis-backed per-host token bucket for outbound HTTP calls.
+
+Shared across all Celery workers so that every task hitting, e.g.,
+``db.satnogs.org`` consumes from the same budget regardless of which worker
+picked it up. Two concurrent windows: per-minute (always) and optional per-hour,
+because Space-Track publishes 20/min *and* 200/hr ceilings.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from urllib.parse import urlparse
+
+from django.conf import settings
+from django.core.cache import cache
+
+from .conf import cache_key
+
+logger = logging.getLogger(__name__)
+
+
+def _bucket_key_minute(host: str) -> str:
+    return cache_key("host_rl", host)
+
+
+def _bucket_key_hour(host: str) -> str:
+    return cache_key("host_rl", "hour", host)
+
+
+def acquire(
+    host: str,
+    per_minute: int,
+    *,
+    per_hour: int = 0,
+    wait_seconds: float = 30.0,
+) -> None:
+    """Block (up to wait_seconds) until tokens are available for *host*."""
+    if not host or per_minute <= 0:
+        return
+
+    # Development has no shared Redis and no traffic worth pacing; blocking here
+    # would only make the dev server feel broken.
+    if not getattr(settings, "PRODUCTION", False):
+        logger.debug(
+            "host rate limit acquire: %s (per_minute=%d, per_hour=%d)",
+            host, per_minute, per_hour,
+        )
+        return
+
+    deadline = time.monotonic() + wait_seconds
+    if per_hour > 0:
+        _wait_for_window(_bucket_key_hour(host), per_hour, timeout=3600, deadline=deadline)
+    _wait_for_window(_bucket_key_minute(host), per_minute, timeout=60, deadline=deadline)
+
+
+def _wait_for_window(key: str, limit: int, *, timeout: int, deadline: float) -> None:
+    while True:
+        count = cache.get(key, 0)
+        if count < limit:
+            try:
+                count = cache.incr(key)
+            except ValueError:
+                cache.set(key, 1, timeout=timeout)
+                count = 1
+            if count <= limit:
+                return
+        sleep_for = min(1.0, max(0.05, deadline - time.monotonic()))
+        if time.monotonic() >= deadline:
+            # Proceeding over budget beats stalling a worker indefinitely; the
+            # upstream's own 429 handling in http.get() is the backstop.
+            logger.warning(
+                "host rate limit wait exceeded for %s (limit=%d, window=%ds)",
+                key, limit, timeout,
+            )
+            return
+        time.sleep(sleep_for)
+
+
+def host_from_url(url: str) -> str:
+    return (urlparse(url).hostname or "").lower()