feat(api_core): allow setting Retry deadline as strict

If a deadline is set as strict, Retry will shorten the last sleep
period to end at the given deadline, and not possibly stretch beyond
it.

Author: plamut

api_core/google/api_core/retry.py

@@ -141,7 +141,9 @@ def exponential_sleep_generator(initial, maximum, multiplier=_DEFAULT_DELAY_MULT
         delay = delay * multiplier
 
 
-def retry_target(target, predicate, sleep_generator, deadline, on_error=None):
+def retry_target(
+    target, predicate, sleep_generator, deadline, on_error=None, strict_deadline=False
+):
     """Call a function and retry if it fails.
 
     This is the lowest-level retry helper. Generally, you'll use the
@@ -156,9 +158,12 @@ def retry_target(target, predicate, sleep_generator, deadline, on_error=None):
         sleep_generator (Iterable[float]): An infinite iterator that determines
             how long to sleep between retries.
         deadline (float): How long to keep retrying the target.
-        on_error (Callable): A function to call while processing a retryable
-            exception.  Any error raised by this function will *not* be
-            caught.
+        on_error (Callable[Exception]): A function to call while processing a
+            retryable exception.  Any error raised by this function will *not*
+            be caught.
+        strict_deadline (bool): If :data:`True`, the last retry will run at
+            ``deadline``, shortening the last sleep interval as necessary.
+            Defaults to :data:`False`.
 
     Returns:
         Any: the return value of the target function.
@@ -191,16 +196,21 @@ def retry_target(target, predicate, sleep_generator, deadline, on_error=None):
                 on_error(exc)
 
         now = datetime_helpers.utcnow()
-        if deadline_datetime is not None and deadline_datetime < now:
-            six.raise_from(
-                exceptions.RetryError(
-                    "Deadline of {:.1f}s exceeded while calling {}".format(
-                        deadline, target
+
+        if deadline_datetime is not None:
+            if deadline_datetime <= now:
+                six.raise_from(
+                    exceptions.RetryError(
+                        "Deadline of {:.1f}s exceeded while calling {}".format(
+                            deadline, target
+                        ),
+                        last_exc,
                     ),
                     last_exc,
-                ),
-                last_exc,
-            )
+                )
+            elif strict_deadline:
+                time_to_deadline = (deadline_datetime - now).total_seconds()
+                sleep = min(time_to_deadline, sleep)
 
         _LOGGER.debug(
             "Retrying due to {}, sleeping {:.1f}s ...".format(last_exc, sleep)
@@ -228,6 +238,9 @@ class Retry(object):
         maximum (float): The maximum amout of time to delay in seconds.
         multiplier (float): The multiplier applied to the delay.
         deadline (float): How long to keep retrying in seconds.
+        strict_deadline (bool): If :data:`True`, the last retry will run at
+            ``deadline``, shortening the last sleep interval as necessary.
+            Defaults to :data:`False`.
     """
 
     def __init__(
@@ -238,21 +251,23 @@ def __init__(
         multiplier=_DEFAULT_DELAY_MULTIPLIER,
         deadline=_DEFAULT_DEADLINE,
         on_error=None,
+        strict_deadline=False,
     ):
         self._predicate = predicate
         self._initial = initial
         self._multiplier = multiplier
         self._maximum = maximum
         self._deadline = deadline
         self._on_error = on_error
+        self._strict_deadline = strict_deadline
 
     def __call__(self, func, on_error=None):
         """Wrap a callable with retry behavior.
 
         Args:
             func (Callable): The callable to add retry behavior to.
-            on_error (Callable): A function to call while processing a
-                retryable exception.  Any error raised by this function will
+            on_error (Callable[Exception]): A function to call while processing
+                a retryable exception. Any error raised by this function will
                 *not* be caught.
 
         Returns:
@@ -275,15 +290,19 @@ def retry_wrapped_func(*args, **kwargs):
                 sleep_generator,
                 self._deadline,
                 on_error=on_error,
+                strict_deadline=self._strict_deadline
             )
 
         return retry_wrapped_func
 
-    def with_deadline(self, deadline):
+    def with_deadline(self, deadline, strict_deadline=False):
         """Return a copy of this retry with the given deadline.
 
         Args:
             deadline (float): How long to keep retrying.
+            strict_deadline (bool): If :data:`True`, the last retry will run at
+                ``deadline``, shortening the last sleep interval as necessary.
+                Defaults to :data:`False`.
 
         Returns:
             Retry: A new retry instance with the given deadline.
@@ -295,6 +314,7 @@ def with_deadline(self, deadline):
             multiplier=self._multiplier,
             deadline=deadline,
             on_error=self._on_error,
+            strict_deadline=strict_deadline,
         )
 
     def with_predicate(self, predicate):
@@ -314,6 +334,7 @@ def with_predicate(self, predicate):
             multiplier=self._multiplier,
             deadline=self._deadline,
             on_error=self._on_error,
+            strict_deadline=self._strict_deadline,
         )
 
     def with_delay(self, initial=None, maximum=None, multiplier=None):
@@ -335,17 +356,20 @@ def with_delay(self, initial=None, maximum=None, multiplier=None):
             multiplier=multiplier if maximum is not None else self._multiplier,
             deadline=self._deadline,
             on_error=self._on_error,
+            strict_deadline=self._strict_deadline,
         )
 
     def __str__(self):
         return (
             "<Retry predicate={}, initial={:.1f}, maximum={:.1f}, "
-            "multiplier={:.1f}, deadline={:.1f}, on_error={}>".format(
+            "multiplier={:.1f}, deadline={:.1f}, on_error={}, "
+            "strict_deadline={}>".format(
                 self._predicate,
                 self._initial,
                 self._maximum,
                 self._multiplier,
                 self._deadline,
                 self._on_error,
+                self._strict_deadline,
             )
         )

api_core/tests/unit/test_retry.py

@@ -162,6 +162,7 @@ def test_constructor_defaults(self):
         assert retry_._multiplier == 2
         assert retry_._deadline == 120
         assert retry_._on_error is None
+        assert retry_._strict_deadline is False
 
     def test_constructor_options(self):
         _some_function = mock.Mock()
@@ -173,49 +174,118 @@ def test_constructor_options(self):
             multiplier=3,
             deadline=4,
             on_error=_some_function,
+            strict_deadline=True,
         )
         assert retry_._predicate == mock.sentinel.predicate
         assert retry_._initial == 1
         assert retry_._maximum == 2
         assert retry_._multiplier == 3
         assert retry_._deadline == 4
         assert retry_._on_error is _some_function
+        assert retry_._strict_deadline is True
 
     def test_with_deadline(self):
-        retry_ = retry.Retry()
-        new_retry = retry_.with_deadline(42)
+        retry_ = retry.Retry(
+            predicate=mock.sentinel.predicate,
+            initial=1,
+            maximum=2,
+            multiplier=3,
+            deadline=4,
+            on_error=mock.sentinel.on_error,
+            strict_deadline=True,
+        )
+        new_retry = retry_.with_deadline(42, strict_deadline=True)
         assert retry_ is not new_retry
         assert new_retry._deadline == 42
+        assert new_retry._strict_deadline is True
+
+        # the rest of the attributes should remain the same
+        assert new_retry._predicate is retry_._predicate
+        assert new_retry._initial == retry_._initial
+        assert new_retry._maximum == retry_._maximum
+        assert new_retry._multiplier == retry_._multiplier
+        assert new_retry._on_error is retry_._on_error
 
     def test_with_predicate(self):
-        retry_ = retry.Retry()
+        retry_ = retry.Retry(
+            predicate=mock.sentinel.predicate,
+            initial=1,
+            maximum=2,
+            multiplier=3,
+            deadline=4,
+            on_error=mock.sentinel.on_error,
+            strict_deadline=True,
+        )
         new_retry = retry_.with_predicate(mock.sentinel.predicate)
         assert retry_ is not new_retry
         assert new_retry._predicate == mock.sentinel.predicate
 
+        # the rest of the attributes should remain the same
+        assert new_retry._deadline == retry_._deadline
+        assert new_retry._strict_deadline == retry_._strict_deadline
+        assert new_retry._initial == retry_._initial
+        assert new_retry._maximum == retry_._maximum
+        assert new_retry._multiplier == retry_._multiplier
+        assert new_retry._on_error is retry_._on_error
+
     def test_with_delay_noop(self):
-        retry_ = retry.Retry()
+        retry_ = retry.Retry(
+            predicate=mock.sentinel.predicate,
+            initial=1,
+            maximum=2,
+            multiplier=3,
+            deadline=4,
+            on_error=mock.sentinel.on_error,
+            strict_deadline=True,
+        )
         new_retry = retry_.with_delay()
         assert retry_ is not new_retry
         assert new_retry._initial == retry_._initial
         assert new_retry._maximum == retry_._maximum
         assert new_retry._multiplier == retry_._multiplier
 
     def test_with_delay(self):
-        retry_ = retry.Retry()
+        retry_ = retry.Retry(
+            predicate=mock.sentinel.predicate,
+            initial=1,
+            maximum=2,
+            multiplier=3,
+            deadline=4,
+            on_error=mock.sentinel.on_error,
+            strict_deadline=True,
+        )
         new_retry = retry_.with_delay(initial=1, maximum=2, multiplier=3)
         assert retry_ is not new_retry
         assert new_retry._initial == 1
         assert new_retry._maximum == 2
         assert new_retry._multiplier == 3
 
+        # the rest of the attributes should remain the same
+        assert new_retry._deadline == retry_._deadline
+        assert new_retry._strict_deadline == retry_._strict_deadline
+        assert new_retry._predicate is retry_._predicate
+        assert new_retry._on_error is retry_._on_error
+
     def test___str__(self):
-        retry_ = retry.Retry()
+        def if_exception_type(exc):
+            return bool(exc)  # the exact logic is irrelevant
+
+        # Explicitly set all attributes as changed Retry defaults should not
+        # cause this test to start failing.
+        retry_ = retry.Retry(
+            predicate=if_exception_type,
+            initial=1.0,
+            maximum=60.0,
+            multiplier=2.0,
+            deadline=120.0,
+            on_error=None,
+            strict_deadline=False,
+        )
         assert re.match(
             (
                 r"<Retry predicate=<function.*?if_exception_type.*?>, "
                 r"initial=1.0, maximum=60.0, multiplier=2.0, deadline=120.0, "
-                r"on_error=None>"
+                r"on_error=None, strict_deadline=False>"
             ),
             str(retry_),
         )
@@ -259,6 +329,55 @@ def test___call___and_execute_retry(self, sleep, uniform):
         sleep.assert_called_once_with(retry_._initial)
         assert on_error.call_count == 1
 
+    # Make uniform return half of its maximum, which is the calculated sleep time.
+    @mock.patch("random.uniform", autospec=True, side_effect=lambda m, n: n / 2.0)
+    @mock.patch("time.sleep", autospec=True)
+    def test___call___and_execute_retry_strict_deadline(self, sleep, uniform):
+
+        on_error = mock.Mock(spec=["__call__"], side_effect=[None] * 10)
+        retry_ = retry.Retry(
+            predicate=retry.if_exception_type(ValueError),
+            initial=1.0,
+            maximum=1024.0,
+            multiplier=2.0,
+            deadline=9.9,
+            strict_deadline=True,
+        )
+
+        utcnow = datetime.datetime.utcnow()
+        utcnow_patcher = mock.patch(
+            "google.api_core.datetime_helpers.utcnow", return_value=utcnow
+        )
+
+        target = mock.Mock(spec=["__call__"], side_effect=[ValueError()] * 10)
+        # __name__ is needed by functools.partial.
+        target.__name__ = "target"
+
+        decorated = retry_(target, on_error=on_error)
+        target.assert_not_called()
+
+        with utcnow_patcher as patched_utcnow:
+            # Make sure that calls to fake time.sleep() also advance the mocked
+            # time clock.
+            def increase_time(sleep_delay):
+                patched_utcnow.return_value += datetime.timedelta(seconds=sleep_delay)
+            sleep.side_effect = increase_time
+
+            with pytest.raises(exceptions.RetryError):
+                decorated("meep")
+
+        assert target.call_count == 5
+        target.assert_has_calls([mock.call("meep")] * 5)
+        assert on_error.call_count == 5
+
+        # check the delays
+        assert sleep.call_count == 4  # once between each successive target calls
+        last_wait = sleep.call_args.args[0]
+        total_wait = sum(call_args.args[0] for call_args in sleep.call_args_list)
+
+        assert last_wait == 2.9   # and not 8.0, because the last delay was shortened
+        assert total_wait == 9.9  # the same as the (strict) deadline
+
     @mock.patch("time.sleep", autospec=True)
     def test___init___without_retry_executed(self, sleep):
         _some_function = mock.Mock()