Skip to content

Commit

Permalink
fix: resolve the issue where rpc timeout of 0 is used when timeout ex…
Browse files Browse the repository at this point in the history
…pires
  • Loading branch information
parthea committed Jan 12, 2025
1 parent 9182916 commit 95cde20
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 7 deletions.
19 changes: 15 additions & 4 deletions google/api_core/timeout.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand All @@ -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)

Expand Down
6 changes: 3 additions & 3 deletions tests/unit/test_timeout.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down

0 comments on commit 95cde20

Please sign in to comment.