diff --git a/src/govoplan_mail/backend/sending/rate_limit.py b/src/govoplan_mail/backend/sending/rate_limit.py index 2cd34c7..8e05ff2 100644 --- a/src/govoplan_mail/backend/sending/rate_limit.py +++ b/src/govoplan_mail/backend/sending/rate_limit.py @@ -1,6 +1,7 @@ from __future__ import annotations import time +import threading from dataclasses import dataclass from redis import Redis @@ -17,6 +18,10 @@ class RateLimitDecision: 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) @@ -30,15 +35,18 @@ def wait_for_rate_limit(*, key: str, messages_per_minute: int, enabled: bool = T 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 Redis - is only used when Celery/worker mode is explicitly enabled. If Redis is - unavailable, it falls back to no distributed wait. + 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 or not _distributed_rate_limit_enabled(): + 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" @@ -56,7 +64,17 @@ def wait_for_rate_limit(*, key: str, messages_per_minute: int, enabled: bool = T now = time.time() client.set(redis_key, now + gap, ex=max(60, int(gap * 10))) except (RedisError, TimeoutError, ValueError): - # Development fallback: do not fail sending because Redis is absent. - waited = 0.0 + 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 diff --git a/tests/test_rate_limit.py b/tests/test_rate_limit.py new file mode 100644 index 0000000..64a8902 --- /dev/null +++ b/tests/test_rate_limit.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import unittest +from unittest.mock import patch + +from redis.exceptions import RedisError + +from govoplan_mail.backend.sending import rate_limit + + +class RateLimitTests(unittest.TestCase): + def setUp(self) -> None: + rate_limit._local_next_allowed.clear() + + def test_local_fallback_waits_between_sends(self) -> None: + now = [100.0] + sleeps: list[float] = [] + + def fake_time() -> float: + return now[0] + + def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + now[0] += seconds + + with ( + patch.object(rate_limit, "_distributed_rate_limit_enabled", return_value=False), + patch.object(rate_limit.time, "time", fake_time), + patch.object(rate_limit.time, "sleep", fake_sleep), + ): + first = rate_limit.wait_for_rate_limit(key="tenant:t:campaign:c", messages_per_minute=4) + second = rate_limit.wait_for_rate_limit(key="tenant:t:campaign:c", messages_per_minute=4) + + self.assertEqual(first.waited_seconds, 0.0) + self.assertEqual(second.waited_seconds, 15.0) + self.assertEqual(sleeps, [15.0]) + + def test_disabled_rate_limit_does_not_reserve_local_slot(self) -> None: + now = [100.0] + sleeps: list[float] = [] + + with ( + patch.object(rate_limit, "_distributed_rate_limit_enabled", return_value=False), + patch.object(rate_limit.time, "time", lambda: now[0]), + patch.object(rate_limit.time, "sleep", lambda seconds: sleeps.append(seconds)), + ): + disabled = rate_limit.wait_for_rate_limit(key="k", messages_per_minute=4, enabled=False) + enabled = rate_limit.wait_for_rate_limit(key="k", messages_per_minute=4) + + self.assertEqual(disabled.waited_seconds, 0.0) + self.assertEqual(enabled.waited_seconds, 0.0) + self.assertEqual(sleeps, []) + + def test_redis_failure_falls_back_to_local_wait(self) -> None: + now = [100.0] + sleeps: list[float] = [] + + def fail_redis_client(): + raise RedisError("redis unavailable") + + def fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + now[0] += seconds + + with ( + patch.object(rate_limit, "_distributed_rate_limit_enabled", return_value=True), + patch.object(rate_limit, "_redis_client", fail_redis_client), + patch.object(rate_limit.time, "time", lambda: now[0]), + patch.object(rate_limit.time, "sleep", fake_sleep), + ): + first = rate_limit.wait_for_rate_limit(key="k", messages_per_minute=60) + second = rate_limit.wait_for_rate_limit(key="k", messages_per_minute=60) + + self.assertEqual(first.waited_seconds, 0.0) + self.assertEqual(second.waited_seconds, 1.0) + self.assertEqual(sleeps, [1.0]) + + +if __name__ == "__main__": + unittest.main()