fix: enforce total timeout in APIBudget._do_acquire instead of per-attempt

airbyte_cdk/sources/streams/call_rate.py

@@ -622,11 +622,12 @@ def _do_acquire(
         :param request: the API request
         :param policy: the matching rate-limiting policy
         :param block: indicates whether to block until a call credit is available
-        :param timeout: maximum time to wait if blocking
+        :param timeout: maximum total time to wait if blocking, in seconds
         :raises: CallRateLimitHit if unable to acquire a call credit
         """
         last_exception = None
         endpoint = self._extract_endpoint(request)
+        deadline = time.monotonic() + timeout if (block and timeout is not None) else None
         # sometimes we spend all budget before a second attempt, so we have a few more attempts
         for attempt in range(1, self._maximum_attempts_to_acquire):
             try:
@@ -635,25 +636,31 @@ def _do_acquire(
                 return
             except CallRateLimitHit as exc:
                 last_exception = exc
-                if block:
-                    if timeout is not None:
-                        time_to_wait = min(timedelta(seconds=timeout), exc.time_to_wait)
-                    else:
-                        time_to_wait = exc.time_to_wait
-                    # Ensure we never sleep for a negative duration.
-                    time_to_wait = max(timedelta(0), time_to_wait)
-                    logger.debug(
-                        f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}). "
-                        f"Sleeping for {time_to_wait} on attempt {attempt}."
-                    )
-                    time.sleep(time_to_wait.total_seconds())
-                else:
+                if not block:
                     logger.debug(
                         f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}) "
                         f"and blocking is disabled."
                     )
                     raise
 
+                time_to_wait = exc.time_to_wait
+                if deadline is not None:
+                    remaining = timedelta(seconds=deadline - time.monotonic())
+                    if remaining <= timedelta(0):
+                        logger.debug(
+                            f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}) "
+                            f"and the total timeout of {timeout}s has been exhausted."
+                        )
+                        raise
+                    time_to_wait = min(time_to_wait, remaining)
+                # Ensure we never sleep for a negative duration.
+                time_to_wait = max(timedelta(0), time_to_wait)
+                logger.debug(
+                    f"Policy {policy} reached call limit for endpoint {endpoint} ({exc.rate}). "
+                    f"Sleeping for {time_to_wait} on attempt {attempt}."
+                )
+                time.sleep(time_to_wait.total_seconds())
+
         if last_exception:
             logger.debug(
                 f"Exhausted all {self._maximum_attempts_to_acquire} attempts to acquire a call for endpoint {endpoint} "

unit_tests/sources/streams/test_call_rate.py

@@ -13,6 +13,7 @@
 
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams.call_rate import (
+    AbstractCallRatePolicy,
     APIBudget,
     CallRateLimitHit,
     FixedWindowCallRatePolicy,
@@ -60,6 +61,117 @@ def enable_cache_fixture():
         os.environ[ENV_REQUEST_CACHE_PATH] = prev_cache_path
 
 
+class StubCallRatePolicy(AbstractCallRatePolicy):
+    def __init__(self, failures: Optional[int] = None):
+        self._failures = failures
+        self.try_acquire_calls = 0
+
+    def matches(self, request: Any) -> bool:
+        return True
+
+    def try_acquire(self, request: Any, weight: int) -> None:
+        self.try_acquire_calls += 1
+        if self._failures is None or self._failures > 0:
+            if self._failures is not None:
+                self._failures -= 1
+            raise CallRateLimitHit(
+                error="x",
+                item=request,
+                weight=1,
+                rate="r",
+                time_to_wait=timedelta(seconds=60),
+            )
+
+    def update(self, available_calls: Optional[int], call_reset_ts: Optional[datetime]) -> None:
+        pass
+
+
+@pytest.fixture
+def fake_clock(mocker):
+    class FakeClock:
+        def __init__(self):
+            self.now = 0.0
+            self.sleep_calls = []
+
+        def monotonic(self) -> float:
+            return self.now
+
+        def sleep(self, seconds: float) -> None:
+            self.sleep_calls.append(seconds)
+            self.now += seconds
+
+    clock = FakeClock()
+    mocker.patch(
+        "airbyte_cdk.sources.streams.call_rate.time.monotonic",
+        side_effect=clock.monotonic,
+    )
+    mocker.patch(
+        "airbyte_cdk.sources.streams.call_rate.time.sleep",
+        side_effect=clock.sleep,
+    )
+    return clock
+
+
+class TestAPIBudgetAcquireTimeout:
+    def test_timeout_is_total_budget(self, fake_clock):
+        policy = StubCallRatePolicy()
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(object(), block=True, timeout=1)
+
+        assert fake_clock.sleep_calls == [1.0]
+        assert policy.try_acquire_calls == 2
+
+    def test_zero_timeout_does_not_sleep(self, fake_clock):
+        policy = StubCallRatePolicy()
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(object(), block=True, timeout=0)
+
+        assert fake_clock.sleep_calls == []
+        assert policy.try_acquire_calls == 1
+
+    def test_succeeds_before_timeout(self, fake_clock):
+        policy = StubCallRatePolicy(failures=2)
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        budget.acquire_call(object(), block=True, timeout=200)
+
+        assert fake_clock.sleep_calls == [60.0, 60.0]
+        assert policy.try_acquire_calls == 3
+
+    def test_none_timeout_preserves_success_behavior(self, fake_clock):
+        policy = StubCallRatePolicy(failures=2)
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        budget.acquire_call(object(), block=True, timeout=None)
+
+        assert fake_clock.sleep_calls == [60.0, 60.0]
+        assert policy.try_acquire_calls == 3
+
+    def test_none_timeout_preserves_attempt_behavior(self, fake_clock):
+        policy = StubCallRatePolicy()
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(object(), block=True, timeout=None)
+
+        assert fake_clock.sleep_calls == [60.0, 60.0, 60.0]
+        assert policy.try_acquire_calls == 3
+
+    def test_non_blocking_raises_immediately(self, fake_clock):
+        policy = StubCallRatePolicy()
+        budget = APIBudget(policies=[policy], maximum_attempts_to_acquire=4)
+
+        with pytest.raises(CallRateLimitHit):
+            budget.acquire_call(object(), block=False, timeout=1)
+
+        assert fake_clock.sleep_calls == []
+        assert policy.try_acquire_calls == 1
+
+
 class TestHttpRequestMatcher:
     try_all_types_of_requests = pytest.mark.parametrize(
         "request_factory",