fix(call-rate): honor RateLimit-Remaining when a reset header is present

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,37 +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:
-                    now: int = TimeClock().now()  # type: ignore[no-untyped-call]
-                    items_to_add = self._get_available_calls(now)
-                    if items_to_add > 0:
-                        added = self._bucket.put(
-                            RateItem(name="dummy", timestamp=now, weight=items_to_add)
-                        )
-                        if not added:
-                            logger.warning(
-                                "Failed to sync rate limit bucket with API: could not add %s items",
-                                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()
-
-    def _get_available_calls(self, now: int) -> int:
-        """Return the number of calls still allowed by the strictest configured rate at the given timestamp (ms)."""
-        return min(
-            rate.limit
-            - sum(1 for item in self._bucket.items if item.timestamp >= now - rate.interval)
-            for rate in self._bucket.rates
-        )
+        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."""

unit_tests/sources/streams/test_call_rate.py

@@ -17,6 +17,7 @@
     APIBudget,
     CallRateLimitHit,
     FixedWindowCallRatePolicy,
+    HttpAPIBudget,
     HttpRequestMatcher,
     HttpRequestRegexMatcher,
     MovingWindowCallRatePolicy,
@@ -306,6 +307,91 @@ 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)
+
     def test_update_available_calls_zero_exhausts_bucket(self):
         with freeze_time("2024-01-01 00:00:00") as frozen:
             policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
@@ -352,7 +438,7 @@ def test_update_available_calls_zero_ignores_expired_entries(self):
             with pytest.raises(CallRateLimitHit):
                 policy.try_acquire("call", weight=1)
 
-    def test_update_available_calls_zero_multiple_rates(self):
+    def test_update_skipped_when_over_cap_rate_blocks_fill(self):
         with freeze_time("2024-01-01 00:00:00"):
             policy = MovingWindowCallRatePolicy(
                 rates=[
@@ -366,21 +452,68 @@ def test_update_available_calls_zero_multiple_rates(self):
 
             policy.update(available_calls=0, call_reset_ts=None)
 
-            assert policy._bucket.count() == 2
-            with pytest.raises(CallRateLimitHit):
-                policy.try_acquire("call", weight=1)
+            # The 1-hour rate exceeds MAX_HEADER_DRIVEN_WAIT, so it is excluded from
+            # _calls_left, but put() still validates the fill against it. The fill is
+            # rejected when it would exhaust that rate, leaving the policy unchanged.
+            assert policy._bucket.count() == 1
 
-    def test_update_positive_available_calls_is_noop(self):
+    def test_update_does_not_raise_allowance(self):
         with freeze_time("2024-01-01 00:00:00"):
             policy = MovingWindowCallRatePolicy(rates=[Rate(10, timedelta(minutes=1))], matchers=[])
             for _ in range(3):
                 policy.try_acquire("call", weight=1)
 
-            policy.update(available_calls=5, call_reset_ts=None)
+            policy.update(available_calls=50, call_reset_ts=None)
             assert policy._bucket.count() == 3
 
-            policy.update(available_calls=0, call_reset_ts=datetime.now())
-            assert policy._bucket.count() == 3
+
+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: