diff --git a/requests_ratelimiter/requests_ratelimiter.py b/requests_ratelimiter/requests_ratelimiter.py index bce8028..fa3f75b 100644 --- a/requests_ratelimiter/requests_ratelimiter.py +++ b/requests_ratelimiter/requests_ratelimiter.py @@ -7,6 +7,7 @@ from pyrate_limiter import Duration, InMemoryBucket, Limiter, Rate from pyrate_limiter.abstracts import AbstractBucket, RateItem +from pyrate_limiter.limiter import combined_lock from requests import PreparedRequest, Response, Session, exceptions from requests.adapters import HTTPAdapter @@ -146,16 +147,17 @@ def _fill_bucket(self, request: PreparedRequest): logger.warning('No buckets available to fill') return - now = bucket.now() - # Use smallest rate interval (first after sorting) rate = sorted(bucket.rates, key=lambda r: r.interval)[0] # Add filler items to saturate the smallest rate limit # This ensures we delay before the next request - for _ in range(rate.limit): - filler_item = RateItem(bucket_name, now, weight=1) - bucket.put(filler_item) + # Each put() resets bucket.failing_rate, which try_acquire() reads under this lock + with combined_lock(self.limiter.lock, blocking=True): + now = bucket.now() + for _ in range(rate.limit): + filler_item = RateItem(bucket_name, now, weight=1) + bucket.put(filler_item) def close(self) -> None: """Close the session or adapter and release all rate-limiter resources. diff --git a/test/test_requests_ratelimiter.py b/test/test_requests_ratelimiter.py index 35fcb3b..a179c74 100644 --- a/test/test_requests_ratelimiter.py +++ b/test/test_requests_ratelimiter.py @@ -4,6 +4,7 @@ """ import pickle +from threading import Thread from time import sleep from unittest.mock import MagicMock, patch @@ -354,6 +355,25 @@ def test_fill_bucket_no_bucket_logs_warning(caplog): assert 'No buckets available' in caplog.text +def test_fill_bucket_waits_for_limiter_lock(): + """A filler put() resets bucket.failing_rate, so it must not run while another thread holds + the limiter lock between its own put() and waiting() + """ + session = get_mock_session(per_second=5) + req = PreparedRequest() + req.url = MOCKED_URL + bucket = session.limiter.bucket_factory[session._bucket_name(req)] + + filler = Thread(target=session._fill_bucket, args=(req,)) + with session.limiter.lock: + filler.start() + filler.join(timeout=0.1) + assert bucket.count() == 0 + + filler.join() + assert bucket.count() == 5 + + @patch_sleep def test_custom_bucket_class(mock_sleep): class MyBucket(InMemoryBucket):