Add google.api.core.retry.Retry decorator

core/google/api/core/retry.py

@@ -12,9 +12,52 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Helpers for retrying functions with exponential back-off."""
+"""Helpers for retrying functions with exponential back-off.
+
+The :cls`Retry` decorator can be used to retry function that raise exceptions
+using exponential backoff. Because a exponential sleep algorithm is used,
+the retry is limited by a `deadline`. The deadline is the maxmimum amount of
+time a method can block. This is used instead of total number of retries
+because it is difficult to ascertain the amount of time a function can block
+when using total number of retries and exponential backoff.
+
+By default, this decorator will retry transient
+API errors (see :func:`if_transient_error`). For example:
+
+.. code-block:: python
+
+    @retry.Retry()
+    def call_flaky_rpc():
+        return client.flaky_rpc()
+
+    # Will retry flaky_rpc() if it raises transient API errors.
+    result = call_flaky_rpc()
+
+You can pass a custom predicate to retry on different exceptions, such as
+waiting for an eventually consistent item to be available:
+
+.. code-block:: python
+
+    @retry.Retry(predicate=if_exception_type(exceptions.NotFound))
+    def check_if_exists():
+        return client.does_thing_exist()
+
+    is_available = check_if_exists()
+
+Some client library methods apply retry automatically. These methods can accept
+a ``retry`` parameter that allows you to configure the behavior:
+
+.. code-block:: python
+
+    my_retry = retry.Retry(deadline=60)
+    result = client.some_method(retry=my_retry)
+
+"""
+
+from __future__ import unicode_literals
 
 import datetime
+import functools
 import logging
 import random
 import time
@@ -38,10 +81,10 @@ def if_exception_type(*exception_types):
         Callable[Exception]: A predicate that returns True if the provided
             exception is of the given type(s).
     """
-    def inner(exception):
+    def if_exception_type_predicate(exception):
         """Bound predicate for checking an exception type."""
         return isinstance(exception, exception_types)
-    return inner
+    return if_exception_type_predicate
 
 
 # pylint: disable=invalid-name
@@ -146,3 +189,129 @@ def retry_target(target, predicate, sleep_generator, deadline):
         time.sleep(sleep)
 
     raise ValueError('Sleep generator stopped yielding sleep values.')
+
+
+@six.python_2_unicode_compatible
+class Retry(object):
+    """Exponential retry decorator.
+
+    This class is a decorator used to add exponential back-off retry behavior
+    to an RPC call.
+
+    Although the default behavior is to retry transient API errors, a
+    different predicate can be provided to retry other exceptions.
+
+    Args:
+        predicate (Callable[Exception]): A callable that should return ``True``
+            if the given exception is retryable.
+        initial (float): The minimum about of time to delay. This must
+            be greater than 0.
+        maximum (float): The maximum about of time to delay.
+        multiplier (float): The multiplier applied to the delay.
+        jitter (float): The maximum about of randomness to apply to the delay.
+        deadline (float): How long to keep retrying.
+    """
+    def __init__(
+            self,
+            predicate=if_transient_error,
+            initial=1,
+            maximum=60,
+            multiplier=2,
+            jitter=_DEFAULT_MAX_JITTER,
+            deadline=60 * 2):
+        self._predicate = predicate
+        self._initial = initial
+        self._multiplier = multiplier
+        self._maximum = maximum
+        self._jitter = jitter
+        self._deadline = deadline
+
+    def __call__(self, func):
+        """Wrap a callable with retry behavior.
+
+        Args:
+            func (Callable): The callable to add retry behavior to.
+
+        Returns:
+            Callable: A callable that will invoke ``func`` with retry
+                behavior.
+        """
+        @six.wraps(func)
+        def retry_wrapped_func(*args, **kwargs):
+            """A wrapper that calls target function with retry."""
+            target = functools.partial(func, *args, **kwargs)
+            sleep_generator = exponential_sleep_generator(
+                self._initial, self._maximum,
+                multiplier=self._multiplier, jitter=self._jitter)
+            return retry_target(
+                target,
+                self._predicate,
+                sleep_generator,
+                self._deadline)
+
+        return retry_wrapped_func
+
+    def with_deadline(self, deadline):
+        """Returns a copy of this retry with the given deadline.
+
+        Args:
+            deadline (float): How long to keep retrying.
+
+        Returns:
+            Retry: A new retry instance with the given deadline.
+        """
+        return Retry(
+            predicate=self._predicate,
+            initial=self._initial,
+            maximum=self._maximum,
+            multiplier=self._multiplier,
+            jitter=self._jitter,
+            deadline=deadline)
+
+    def with_predicate(self, predicate):
+        """Returns a copy of this retry with the given predicate.
+
+        Args:
+            predicate (Callable[Exception]): A callable that should return
+                ``True`` if the given exception is retryable.
+
+        Returns:
+            Retry: A new retry instance with the given predicate.
+        """
+        return Retry(
+            predicate=predicate,
+            initial=self._initial,
+            maximum=self._maximum,
+            multiplier=self._multiplier,
+            jitter=self._jitter,
+            deadline=self._deadline)
+
+    def with_delay(
+            self, initial=None, maximum=None, multiplier=None, jitter=None):
+        """Returns a copy of this retry with the given delay options.
+
+        Args:
+            initial (float): The minimum about of time to delay. This must
+                be greater than 0.
+            maximum (float): The maximum about of time to delay.
+            multiplier (float): The multiplier applied to the delay.
+            jitter (float): The maximum about of randomness to apply to the
+                delay.
+
+        Returns:
+            Retry: A new retry instance with the given predicate.
+        """
+        return Retry(
+            predicate=self._predicate,
+            initial=initial if initial is not None else self._initial,
+            maximum=maximum if maximum is not None else self._maximum,
+            multiplier=multiplier if maximum is not None else self._multiplier,
+            jitter=jitter if jitter is not None else self._jitter,
+            deadline=self._deadline)
+
+    def __str__(self):
+        return (
+            '<Retry predicate={}, initial={:.1f}, maximum={:.1f}, '
+            'multiplier={:.1f}, jitter={:.1f}, deadline={:.1f}>'.format(
+                self._predicate, self._initial, self._maximum,
+                self._multiplier, self._jitter, self._deadline))

core/tests/unit/api_core/test_retry.py

@@ -14,6 +14,7 @@
 
 import datetime
 import itertools
+import re
 
 import mock
 import pytest
@@ -45,27 +46,28 @@ def test_if_transient_error():
 
 def test_exponential_sleep_generator_base_2():
     gen = retry.exponential_sleep_generator(
-        1, 60, 2, jitter=0.0)
+        1, 60, multiplier=2, jitter=0.0)
 
     result = list(itertools.islice(gen, 8))
     assert result == [1, 2, 4, 8, 16, 32, 60, 60]
 
 
-@mock.patch('random.uniform')
+@mock.patch('random.uniform', autospec=True)
 def test_exponential_sleep_generator_jitter(uniform):
     uniform.return_value = 1
     gen = retry.exponential_sleep_generator(
-        1, 60, 2, jitter=2.2)
+        1, 60, multiplier=2, jitter=2.2)
 
     result = list(itertools.islice(gen, 7))
     assert result == [1, 3, 7, 15, 31, 60, 60]
     uniform.assert_called_with(0.0, 2.2)
 
 
-@mock.patch('time.sleep')
+@mock.patch('time.sleep', autospec=True)
 @mock.patch(
     'google.api.core.helpers.datetime_helpers.utcnow',
-    return_value=datetime.datetime.min)
+    return_value=datetime.datetime.min,
+    autospec=True)
 def test_retry_target_success(utcnow, sleep):
     predicate = retry.if_exception_type(ValueError)
     call_count = [0]
@@ -83,10 +85,11 @@ def target():
     sleep.assert_has_calls([mock.call(0), mock.call(1)])
 
 
-@mock.patch('time.sleep')
+@mock.patch('time.sleep', autospec=True)
 @mock.patch(
     'google.api.core.helpers.datetime_helpers.utcnow',
-    return_value=datetime.datetime.min)
+    return_value=datetime.datetime.min,
+    autospec=True)
 def test_retry_target_non_retryable_error(utcnow, sleep):
     predicate = retry.if_exception_type(ValueError)
     exception = TypeError()
@@ -99,9 +102,9 @@ def test_retry_target_non_retryable_error(utcnow, sleep):
     sleep.assert_not_called()
 
 
-@mock.patch('time.sleep')
+@mock.patch('time.sleep', autospec=True)
 @mock.patch(
-    'google.api.core.helpers.datetime_helpers.utcnow')
+    'google.api.core.helpers.datetime_helpers.utcnow', autospec=True)
 def test_retry_target_deadline_exceeded(utcnow, sleep):
     predicate = retry.if_exception_type(ValueError)
     exception = ValueError('meep')
@@ -127,3 +130,97 @@ def test_retry_target_bad_sleep_generator():
     with pytest.raises(ValueError, match='Sleep generator'):
         retry.retry_target(
             mock.sentinel.target, mock.sentinel.predicate, [], None)
+
+
+class TestRetry(object):
+    def test_constructor_defaults(self):
+        retry_ = retry.Retry()
+        assert retry_._predicate == retry.if_transient_error
+        assert retry_._initial == 1
+        assert retry_._maximum == 60
+        assert retry_._multiplier == 2
+        assert retry_._jitter == retry._DEFAULT_MAX_JITTER
+        assert retry_._deadline == 120
+
+    def test_constructor_options(self):
+        retry_ = retry.Retry(
+            predicate=mock.sentinel.predicate,
+            initial=1,
+            maximum=2,
+            multiplier=3,
+            jitter=4,
+            deadline=5)
+        assert retry_._predicate == mock.sentinel.predicate
+        assert retry_._initial == 1
+        assert retry_._maximum == 2
+        assert retry_._multiplier == 3
+        assert retry_._jitter == 4
+        assert retry_._deadline == 5
+
+    def test_with_deadline(self):
+        retry_ = retry.Retry()
+        new_retry = retry_.with_deadline(42)
+        assert retry_ is not new_retry
+        assert new_retry._deadline == 42
+
+    def test_with_predicate(self):
+        retry_ = retry.Retry()
+        new_retry = retry_.with_predicate(mock.sentinel.predicate)
+        assert retry_ is not new_retry
+        assert new_retry._predicate == mock.sentinel.predicate
+
+    def test_with_delay_noop(self):
+        retry_ = retry.Retry()
+        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
+        assert new_retry._jitter == retry_._jitter
+
+    def test_with_delay(self):
+        retry_ = retry.Retry()
+        new_retry = retry_.with_delay(
+            initial=1, maximum=2, multiplier=3, jitter=4)
+        assert retry_ is not new_retry
+        assert new_retry._initial == 1
+        assert new_retry._maximum == 2
+        assert new_retry._multiplier == 3
+        assert new_retry._jitter == 4
+
+    def test___str__(self):
+        retry_ = retry.Retry()
+        assert re.match((
+            r'<Retry predicate=<function.*?if_exception_type.*?>, '
+            r'initial=1.0, maximum=60.0, multiplier=2.0, jitter=0.2, '
+            r'deadline=120.0>'),
+            str(retry_))
+
+    @mock.patch('time.sleep', autospec=True)
+    def test___call___and_execute_success(self, sleep):
+        retry_ = retry.Retry()
+        target = mock.Mock(spec=['__call__'], return_value=42)
+
+        decorated = retry_(target)
+        target.assert_not_called()
+
+        result = decorated('meep')
+
+        assert result == 42
+        target.assert_called_once_with('meep')
+        sleep.assert_not_called()
+
+    @mock.patch('time.sleep', autospec=True)
+    def test___call___and_execute_retry(self, sleep):
+        retry_ = retry.Retry(predicate=retry.if_exception_type(ValueError))
+        target = mock.Mock(spec=['__call__'], side_effect=[ValueError(), 42])
+
+        decorated = retry_(target)
+        target.assert_not_called()
+
+        result = decorated('meep')
+
+        assert result == 42
+        assert target.call_count == 2
+        target.assert_has_calls([mock.call('meep'), mock.call('meep')])
+        sleep.assert_called_once_with(retry_._initial)