fix(mail): rate-limit direct and fallback sends
This commit is contained in:
@@ -1,6 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import time
|
import time
|
||||||
|
import threading
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
from redis import Redis
|
from redis import Redis
|
||||||
@@ -17,6 +18,10 @@ class RateLimitDecision:
|
|||||||
waited_seconds: float
|
waited_seconds: float
|
||||||
|
|
||||||
|
|
||||||
|
_local_rate_limit_lock = threading.Lock()
|
||||||
|
_local_next_allowed: dict[str, float] = {}
|
||||||
|
|
||||||
|
|
||||||
def _redis_client() -> Redis:
|
def _redis_client() -> Redis:
|
||||||
return Redis.from_url(settings.redis_url, decode_responses=True)
|
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
|
The implementation stores the next allowed send timestamp per key. A Redis
|
||||||
lock keeps multiple Celery processes from reading/updating the timestamp at
|
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
|
the same time. Direct local development runs do not have a broker, so they
|
||||||
is only used when Celery/worker mode is explicitly enabled. If Redis is
|
use a process-local limiter. If Redis is unavailable in worker mode, the
|
||||||
unavailable, it falls back to no distributed wait.
|
process-local fallback still protects single-process sends.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
messages_per_minute = max(1, int(messages_per_minute or 1))
|
messages_per_minute = max(1, int(messages_per_minute or 1))
|
||||||
gap = 60.0 / messages_per_minute
|
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)
|
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"
|
redis_key = f"govoplan:ratelimit:{key}:next_allowed"
|
||||||
lock_key = f"govoplan:ratelimit:{key}:lock"
|
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()
|
now = time.time()
|
||||||
client.set(redis_key, now + gap, ex=max(60, int(gap * 10)))
|
client.set(redis_key, now + gap, ex=max(60, int(gap * 10)))
|
||||||
except (RedisError, TimeoutError, ValueError):
|
except (RedisError, TimeoutError, ValueError):
|
||||||
# Development fallback: do not fail sending because Redis is absent.
|
waited = _wait_for_local_rate_limit(key, gap)
|
||||||
waited = 0.0
|
|
||||||
|
|
||||||
return RateLimitDecision(key=key, messages_per_minute=messages_per_minute, gap_seconds=gap, waited_seconds=waited)
|
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
|
||||||
|
|||||||
80
tests/test_rate_limit.py
Normal file
80
tests/test_rate_limit.py
Normal file
@@ -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()
|
||||||
Reference in New Issue
Block a user