diff --git a/airbyte_cdk/sources/streams/call_rate.py b/airbyte_cdk/sources/streams/call_rate.py index 4a06db3b2..725f90d23 100644 --- a/airbyte_cdk/sources/streams/call_rate.py +++ b/airbyte_cdk/sources/streams/call_rate.py @@ -15,7 +15,7 @@ import requests import requests_cache -from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock +from pyrate_limiter import InMemoryBucket, Limiter, RateItem, TimeClock, binary_search from pyrate_limiter import Rate as PyRateRate from pyrate_limiter.exceptions import BucketFullException @@ -429,6 +429,11 @@ class MovingWindowCallRatePolicy(BaseCallRatePolicy): This strategy requires saving of timestamps of all requests within a window. """ + # Header-driven updates must not induce waits longer than this; sources heartbeat well above it. + # Policies with longer windows are left unchanged rather than parking a worker for the whole window. + MAX_HEADER_DRIVEN_WAIT = timedelta(minutes=10) + _MAX_HEADER_DRIVEN_WAIT_MS = int(MAX_HEADER_DRIVEN_WAIT.total_seconds() * 1000) + def __init__(self, rates: list[Rate], matchers: list[RequestMatcher]): """Constructor @@ -478,22 +483,64 @@ def update( ) -> None: """Adjust call bucket to reflect the state of the API server - :param available_calls: - :param call_reset_ts: + The bucket is filled with dummy calls until what it still allows matches what the API + reports as available. Updates only ever lower the local allowance: when the API reports + more available calls than the configured rates allow, the configured rates win. + + `call_reset_ts` is not used. A moving window has no reset point, so the only actionable + part of the API feedback is the number of calls left; the window length stays the one + the rates were configured with. + + When several rates are configured, the update applies to the most constraining one. A + header describing a coarser window only starts to bite near the end of that window; + mapping headers to a specific rate is deliberately out of scope. + + The subtraction below compares `_calls_left`, which counts bucket entries (weight units; + `put` stores `weight` copies), with `available_calls`, which counts requests. They + coincide for unweighted policies but diverge when weighted matchers are used with a + remaining header. + + :param available_calls: number of calls the API reports as still available + :param call_reset_ts: unused, see above :return: """ - if ( - available_calls is not None and call_reset_ts is None - ): # we do our best to sync buckets with API - if available_calls == 0: - with self._limiter.lock: - items_to_add = self._bucket.count() < self._bucket.rates[0].limit - if items_to_add > 0: - now: int = TimeClock().now() # type: ignore[no-untyped-call] - self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add)) - # TODO: add support if needed, it might be that it is not possible to make a good solution for this case - # if available_calls is not None and call_reset_ts is not None: - # ts = call_reset_ts.timestamp() + if available_calls is None: + return + + available_calls = max(0, available_calls) + with self._limiter.lock: + now: int = TimeClock().now() # type: ignore[no-untyped-call] + calls_left = self._calls_left(now) + if calls_left is None: + return + + items_to_add = calls_left - available_calls + if items_to_add > 0: + logger.debug( + "got rate limit update from api, adjusting available calls from %s to %s", + calls_left, + available_calls, + ) + if not self._bucket.put(RateItem(name="dummy", timestamp=now, weight=items_to_add)): + logger.warning( + "could not adjust available calls: calls_left=%s, available_calls=%s, " + "rejected_weight=%s", + calls_left, + available_calls, + items_to_add, + ) + + def _calls_left(self, now: int) -> Optional[int]: + """Number of calls the bucket still allows, i.e. the most constraining of all rates.""" + items = self._bucket.items + calls_left = [] + for rate in self._bucket.rates: + if rate.interval > self._MAX_HEADER_DRIVEN_WAIT_MS: + continue + lower_bound_idx = binary_search(items, now - rate.interval) + calls_used = len(items) - lower_bound_idx if lower_bound_idx >= 0 else 0 + calls_left.append(rate.limit - calls_used) + return min(calls_left) if calls_left else None def __str__(self) -> str: """Return a human-friendly description of the moving window rate policy for logging purposes.""" diff --git a/unit_tests/sources/streams/test_call_rate.py b/unit_tests/sources/streams/test_call_rate.py index a423fe573..8811b1470 100644 --- a/unit_tests/sources/streams/test_call_rate.py +++ b/unit_tests/sources/streams/test_call_rate.py @@ -16,6 +16,7 @@ APIBudget, CallRateLimitHit, FixedWindowCallRatePolicy, + HttpAPIBudget, HttpRequestMatcher, HttpRequestRegexMatcher, MovingWindowCallRatePolicy, @@ -305,6 +306,140 @@ def test_multiple_limit_rates(self): assert excinfo.value.time_to_wait.total_seconds() == pytest.approx(3600, 0.1) assert str(excinfo.value) == "Bucket for item=call with Rate limit=2/1.0h is already full" + def test_update_available_calls_with_reset_ts(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=2, call_reset_ts=datetime.now()) + + policy.try_acquire("call", weight=1) + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_only_lowers_allowance(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=50, call_reset_ts=datetime.now()) + + for _ in range(10): + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_is_noop_without_available_calls(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=None, call_reset_ts=datetime.now()) + policy.update(available_calls=None, call_reset_ts=None) + + for _ in range(10): + policy.try_acquire("call", weight=1) + + def test_update_respects_the_most_constraining_rate(self): + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(10, timedelta(seconds=1)), + Rate(5, timedelta(minutes=1)), + ], + matchers=[], + ) + + policy.update(available_calls=1, call_reset_ts=None) + + policy.try_acquire("call", weight=1) + with pytest.raises(CallRateLimitHit) as exc: + policy.try_acquire("call", weight=1) + assert exc.value.rate == "limit=5/1.0m" + assert exc.value.time_to_wait.total_seconds() == pytest.approx(60, 0.1) + + def test_update_available_calls_zero_fills_bucket(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + def test_update_ignores_rates_over_header_wait_cap(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[]) + + policy.update(available_calls=0, call_reset_ts=None) + + for _ in range(100): + policy.try_acquire("call", weight=1) + + def test_update_caps_to_eligible_rate(self): + policy = MovingWindowCallRatePolicy( + rates=[ + Rate(10, timedelta(minutes=1)), + Rate(100, timedelta(minutes=15)), + ], + matchers=[], + ) + + policy.update(available_calls=0, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit) as exc: + policy.try_acquire("call", weight=1) + assert exc.value.time_to_wait.total_seconds() <= 600 + + def test_update_clamps_negative_available_calls(self): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + + policy.update(available_calls=-1, call_reset_ts=None) + + with pytest.raises(CallRateLimitHit): + policy.try_acquire("call", weight=1) + + +class TestHttpAPIBudget: + def test_update_from_response(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict( + { + "RateLimit-Remaining": "1", + "RateLimit-Reset": "60", + "RateLimit-Limit": "60", + } + ) + response.status_code = 200 + + budget.update_from_response(request, response) + + budget.acquire_call(request, block=False) + with pytest.raises(CallRateLimitHit): + budget.acquire_call(request, block=False) + + def test_update_from_429_response(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict() + response.status_code = 429 + + budget.update_from_response(request, response) + + with pytest.raises(CallRateLimitHit) as exc: + budget.acquire_call(request, block=False) + assert exc.value.time_to_wait.total_seconds() <= 600 + + def test_update_from_429_response_ignores_over_cap_policy(self, mocker): + policy = MovingWindowCallRatePolicy(rates=[Rate(100, timedelta(minutes=15))], matchers=[]) + budget = HttpAPIBudget(policies=[policy]) + request = Request("GET", "https://example.com") + response = mocker.Mock(spec=requests.Response) + response.headers = requests.structures.CaseInsensitiveDict() + response.status_code = 429 + + budget.update_from_response(request, response) + + for _ in range(100): + budget.acquire_call(request, block=False) + class TestHttpStreamIntegration: def test_without_cache(self, mocker, requests_mock):