81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
"""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()
|