from __future__ import annotations import time import threading from dataclasses import dataclass from redis import Redis from redis.exceptions import RedisError from govoplan_mail.backend.runtime import settings @dataclass(frozen=True, slots=True) class RateLimitDecision: key: str messages_per_minute: int gap_seconds: float waited_seconds: float _local_rate_limit_lock = threading.Lock() _local_next_allowed: dict[str, float] = {} def _redis_client() -> Redis: return Redis.from_url(settings.redis_url, decode_responses=True) def _distributed_rate_limit_enabled() -> bool: return bool(getattr(settings, "celery_enabled", False)) def wait_for_rate_limit(*, key: str, messages_per_minute: int, enabled: bool = True) -> RateLimitDecision: """Throttle sends across worker processes using Redis when available. The implementation stores the next allowed send timestamp per key. A Redis lock keeps multiple Celery processes from reading/updating the timestamp at the same time. Direct local development runs do not have a broker, so they use a process-local limiter. If Redis is unavailable in worker mode, the process-local fallback still protects single-process sends. """ messages_per_minute = max(1, int(messages_per_minute or 1)) gap = 60.0 / messages_per_minute if not enabled: return RateLimitDecision(key=key, messages_per_minute=messages_per_minute, gap_seconds=gap, waited_seconds=0.0) if not _distributed_rate_limit_enabled(): waited = _wait_for_local_rate_limit(key, gap) return RateLimitDecision(key=key, messages_per_minute=messages_per_minute, gap_seconds=gap, waited_seconds=waited) redis_key = f"govoplan:ratelimit:{key}:next_allowed" lock_key = f"govoplan:ratelimit:{key}:lock" waited = 0.0 try: client = _redis_client() with client.lock(lock_key, timeout=30, blocking_timeout=30): now = time.time() raw_next = client.get(redis_key) next_allowed = float(raw_next) if raw_next else now if next_allowed > now: waited = next_allowed - now time.sleep(waited) now = time.time() client.set(redis_key, now + gap, ex=max(60, int(gap * 10))) except (RedisError, TimeoutError, ValueError): waited = _wait_for_local_rate_limit(key, gap) return RateLimitDecision(key=key, messages_per_minute=messages_per_minute, gap_seconds=gap, waited_seconds=waited) def _wait_for_local_rate_limit(key: str, gap: float) -> float: with _local_rate_limit_lock: now = time.time() next_allowed = _local_next_allowed.get(key, now) waited = max(0.0, next_allowed - now) _local_next_allowed[key] = max(now, next_allowed) + gap if waited > 0: time.sleep(waited) return waited