diff --git a/freerelay/middleware/rate_limit.py b/freerelay/middleware/rate_limit.py index 97e07bc..6fba3b7 100644 --- a/freerelay/middleware/rate_limit.py +++ b/freerelay/middleware/rate_limit.py @@ -7,6 +7,7 @@ from __future__ import annotations +import asyncio import logging import time from collections import OrderedDict @@ -44,10 +45,15 @@ def __init__(self, rate: float, capacity: int) -> None: rate: Tokens added per second. capacity: Maximum tokens (burst capacity). """ + if rate <= 0: + raise ValueError("rate must be greater than zero") + if capacity <= 0: + raise ValueError("capacity must be greater than zero") + self.rate = rate self.capacity = capacity self.tokens = float(capacity) - self.last_refill = time.time() + self.last_refill = time.monotonic() def consume(self) -> bool: """ @@ -57,7 +63,7 @@ def consume(self) -> bool: True if a token was available (request allowed). False if no tokens (request should be rate limited). """ - now = time.time() + now = time.monotonic() elapsed = now - self.last_refill self.tokens = min(self.capacity, self.tokens + elapsed * self.rate) self.last_refill = now @@ -83,12 +89,18 @@ def __init__( requests_per_minute: int = 60, burst_capacity: int = 10, ) -> None: + if requests_per_minute < 1: + raise ValueError("requests_per_minute must be at least 1") + if burst_capacity < 1: + raise ValueError("burst_capacity must be at least 1") + super().__init__(app) # type: ignore[arg-type] self.requests_per_minute = requests_per_minute self.burst_capacity = burst_capacity self._buckets: OrderedDict[str, TokenBucket] = OrderedDict() self._rate = requests_per_minute / 60.0 # tokens per second self._request_count = 0 + self._state_lock = asyncio.Lock() def _get_bucket(self, client_ip: str) -> TokenBucket: """Get or create a token bucket for a client with LRU eviction.""" @@ -108,7 +120,7 @@ def _get_bucket(self, client_ip: str) -> TokenBucket: def _cleanup_stale_buckets(self) -> None: """Remove buckets that haven't been used in over 5 minutes.""" - now = time.time() + now = time.monotonic() stale_threshold = 300.0 # 5 minutes stale_keys = [ key @@ -132,17 +144,19 @@ async def dispatch( if any(path.startswith(p) for p in _SKIP_PATHS): return await call_next(request) - # Periodic cleanup of stale buckets - self._request_count += 1 - if self._request_count % _CLEANUP_INTERVAL == 0: - self._cleanup_stale_buckets() + async with self._state_lock: + # Periodic cleanup of stale buckets + self._request_count += 1 + if self._request_count % _CLEANUP_INTERVAL == 0: + self._cleanup_stale_buckets() - # Identify by user_id if available (from AuthMiddleware), otherwise IP - user_id = getattr(request.state, "user_id", None) - client_id = user_id or (request.client.host if request.client else "unknown") - bucket = self._get_bucket(client_id) + # Identify by user_id if available (from AuthMiddleware), otherwise IP + user_id = getattr(request.state, "user_id", None) + client_id = user_id or (request.client.host if request.client else "unknown") + bucket = self._get_bucket(client_id) + allowed = bucket.consume() - if not bucket.consume(): + if not allowed: logger.warning("Rate limit exceeded for %s", client_id) return JSONResponse( status_code=429, diff --git a/tests/unit/test_rate_limit.py b/tests/unit/test_rate_limit.py new file mode 100644 index 0000000..a8cab5a --- /dev/null +++ b/tests/unit/test_rate_limit.py @@ -0,0 +1,43 @@ +from unittest.mock import patch + +import pytest + +from freerelay.middleware.rate_limit import TokenBucket + + +@pytest.mark.parametrize("rate, capacity", [(0, 1), (-1, 1), (1, 0), (1, -1)]) +def test_token_bucket_rejects_non_positive_settings(rate, capacity): + with pytest.raises(ValueError): + TokenBucket(rate, capacity) + + +def test_token_bucket_uses_monotonic_clock_for_refills(): + with patch("freerelay.middleware.rate_limit.time.monotonic", side_effect=[10.0, 10.0, 10.0, 11.0]): + bucket = TokenBucket(rate=1, capacity=1) + assert bucket.consume() + assert not bucket.consume() + + bucket.tokens = 0 + assert bucket.consume() + + +@pytest.mark.asyncio +async def test_concurrent_requests_share_bucket_atomically(): + from unittest.mock import AsyncMock, Mock + + from freerelay.middleware.rate_limit import RateLimitMiddleware + + middleware = RateLimitMiddleware(Mock(), requests_per_minute=1, burst_capacity=1) + request = Mock() + request.url.path = "/v1/chat/completions" + request.state.user_id = "tenant" + request.client.host = "127.0.0.1" + call_next = AsyncMock() + + responses = await __import__("asyncio").gather( + middleware.dispatch(request, call_next), + middleware.dispatch(request, call_next), + ) + + assert sum(response.status_code == 429 for response in responses) == 1 + assert call_next.await_count == 1