From 95cde20cdbf1c7269f1d06cd8554d01cdf6ceab7 Mon Sep 17 00:00:00 2001 From: Anthonios Partheniou Date: Sun, 12 Jan 2025 13:30:06 +0000 Subject: [PATCH] fix: resolve the issue where rpc timeout of 0 is used when timeout expires --- google/api_core/timeout.py | 19 +++++++++++++++---- tests/unit/test_timeout.py | 6 +++--- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/google/api_core/timeout.py b/google/api_core/timeout.py index 868e3e9f..6253b187 100644 --- a/google/api_core/timeout.py +++ b/google/api_core/timeout.py @@ -102,8 +102,7 @@ def __call__(self, func): def func_with_timeout(*args, **kwargs): """Wrapped function that adds timeout.""" - remaining_timeout = self._timeout - if remaining_timeout is not None: + if self._timeout is not None: # All calculations are in seconds now_timestamp = self._clock().timestamp() @@ -114,8 +113,20 @@ def func_with_timeout(*args, **kwargs): now_timestamp = first_attempt_timestamp time_since_first_attempt = now_timestamp - first_attempt_timestamp - # Avoid setting negative timeout - kwargs["timeout"] = max(0, self._timeout - time_since_first_attempt) + remaining_timeout = self._timeout - time_since_first_attempt + + # Although the `deadline` parameter in `google.api_core.retry.Retry` + # is deprecated, and should be treated the same as the `timeout`, + # it is still possible for `deadline` argument in google.api_core.retry.Retry + # to be larger than the `timeout`. + # Avoid setting negative timeout or timeout less than 5 seconds when the `timeout` + # has expired. + # See https://github.com/googleapis/python-api-core/issues/654 + # Revert back to the original timeout when this happens + if remaining_timeout < 5: + remaining_timeout = self._timeout + + kwargs["timeout"] = remaining_timeout return func(*args, **kwargs) diff --git a/tests/unit/test_timeout.py b/tests/unit/test_timeout.py index 60a2e65d..d923bf3a 100644 --- a/tests/unit/test_timeout.py +++ b/tests/unit/test_timeout.py @@ -82,11 +82,11 @@ def _clock(): wrapped() target.assert_called_with(timeout=41.0) wrapped() - target.assert_called_with(timeout=3.0) + target.assert_called_with(timeout=42.0) wrapped() - target.assert_called_with(timeout=0.0) + target.assert_called_with(timeout=42.0) wrapped() - target.assert_called_with(timeout=0.0) + target.assert_called_with(timeout=42.0) def test_apply_no_timeout(self): target = mock.Mock(spec=["__call__", "__name__"], __name__="target")