Add google.api.core.retry with base retry functionality

core/google/api/core/exceptions.py

@@ -40,6 +40,29 @@ class GoogleAPIError(Exception):
     pass
 
 
+@six.python_2_unicode_compatible
+class RetryError(GoogleAPIError):
+    """Raised when a function has exhausted all of its available retries.
+
+    Args:
+        message (str): The exception message.
+        cause (Exception): The last exception raised when retring the
+            function.
+    """
+    def __init__(self, message, cause):
+        super(RetryError, self).__init__(message)
+        self.message = message
+        self._cause = cause
+
+    @property
+    def cause(self):
+        """The last exception raised when retrying the function."""
+        return self._cause
+
+    def __str__(self):
+        return '{}, last exception: {}'.format(self.message, self.cause)
+
+
 class _GoogleAPICallErrorMeta(type):
     """Metaclass for registering GoogleAPICallError subclasses."""
     def __new__(mcs, name, bases, class_dict):

core/google/api/core/helpers/datetime_helpers.py

@@ -0,0 +1,22 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helpers for :mod:`datetime`."""
+
+import datetime
+
+
+def utcnow():
+    """A :meth:`datetime.datetime.utcnow()` alias to allow mocking in tests."""
+    return datetime.datetime.utcnow()

core/google/api/core/retry.py

@@ -0,0 +1,148 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Helpers for retrying functions with exponential back-off."""
+
+import datetime
+import logging
+import random
+import time
+
+import six
+
+from google.api.core import exceptions
+from google.api.core.helpers import datetime_helpers
+
+_LOGGER = logging.getLogger(__name__)
+_DEFAULT_MAX_JITTER = 0.2
+
+
+def if_exception_type(*exception_types):
+    """Creates a predicate to check if the exception is of a given type.
+
+    Args:
+        exception_types (Sequence[type]): The exception types to check for.
+
+    Returns:
+        Callable[Exception]: A predicate that returns True if the provided
+            exception is of the given type(s).
+    """
+    def inner(exception):
+        """Bound predicate for checking an exception type."""
+        return isinstance(exception, exception_types)
+    return inner
+
+
+# pylint: disable=invalid-name
+# Pylint sees this as a constant, but it is also an alias that should be
+# considered a function.
+if_transient_error = if_exception_type((
+    exceptions.InternalServerError,
+    exceptions.TooManyRequests))
+"""A predicate that checks if an exception is a transient API error.
+
+The following server errors are considered transient:
+
+- :class:`google.api.core.exceptions.InternalServerError` - HTTP 500, gRPC
+    ``INTERNAL(13)`` and its subclasses.
+- :class:`google.api.core.exceptions.TooManyRequests` - HTTP 429
+- :class:`google.api.core.exceptions.ResourceExhausted` - gRPC
+    ``RESOURCE_EXHAUSTED(8)``
+"""
+# pylint: enable=invalid-name
+
+
+def exponential_sleep_generator(
+        initial, maximum, multiplier=2, jitter=_DEFAULT_MAX_JITTER):
+    """Generates sleep intervals based on the exponential back-off algorithm.
+
+    This implements the `Truncated Exponential Back-off`_ algorithm.
+
+    .. _Truncated Exponential Back-off:
+        https://cloud.google.com/storage/docs/exponential-backoff
+
+    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.
+
+    Yields:
+        float: successive sleep intervals.
+    """
+    delay = initial
+    while True:
+        yield delay
+        delay = min(
+            delay * multiplier + random.uniform(0, jitter), maximum)
+
+
+def retry_target(target, predicate, sleep_generator, deadline):
+    """Call a function and retry if it fails.
+
+    This is the lowest-level retry helper. Generally, you'll use the
+    higher-level retry helper :class:`Retry`.
+
+    Args:
+        target(Callable): The function to call and retry. This must be a
+            nullary function - apply arguments with `functools.partial`.
+        predicate (Callable[Exception]): A callable used to determine if an
+            exception raised by the target should be considered retryable.
+            It should return True to retry or False otherwise.
+        sleep_generator (Iterator[float]): An infinite iterator that determines
+            how long to sleep between retries.
+        deadline (float): How long to keep retrying the target.
+
+    Returns:
+        Any: the return value of the target function.
+
+    Raises:
+        google.api.core.RetryError: If the deadline is exceeded while retrying.
+        ValueError: If the sleep generator stops yielding values.
+        Exception: If the target raises a method that isn't retryable.
+    """
+    if deadline is not None:
+        deadline_datetime = (
+            datetime_helpers.utcnow() + datetime.timedelta(seconds=deadline))
+    else:
+        deadline_datetime = None
+
+    last_exc = None
+
+    for sleep in sleep_generator:
+        try:
+            return target()
+
+        # pylint: disable=broad-except
+        # This function explicitly must deal with broad exceptions.
+        except Exception as exc:
+            if not predicate(exc):
+                raise
+            last_exc = 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),
+                    last_exc),
+                last_exc)
+
+        _LOGGER.debug('Retrying due to {}, sleeping {:.1f}s ...'.format(
+            last_exc, sleep))
+        time.sleep(sleep)
+
+    raise ValueError('Sleep generator stopped yielding sleep values.')

core/tests/unit/api_core/helpers/test_datetime_helpers.py

@@ -0,0 +1,22 @@
+# Copyright 2017, Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+
+from google.api.core.helpers import datetime_helpers
+
+
+def test_utcnow():
+    result = datetime_helpers.utcnow()
+    assert isinstance(result, datetime.datetime)

core/tests/unit/api_core/test_retry.py

@@ -0,0 +1,129 @@
+# Copyright 2017 Google Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import itertools
+
+import mock
+import pytest
+
+from google.api.core import exceptions
+from google.api.core import retry
+
+
+def test_if_exception_type():
+    predicate = retry.if_exception_type(ValueError)
+
+    assert predicate(ValueError())
+    assert not predicate(TypeError())
+
+
+def test_if_exception_type_multiple():
+    predicate = retry.if_exception_type(ValueError, TypeError)
+
+    assert predicate(ValueError())
+    assert predicate(TypeError())
+    assert not predicate(RuntimeError())
+
+
+def test_if_transient_error():
+    assert retry.if_transient_error(exceptions.InternalServerError(''))
+    assert retry.if_transient_error(exceptions.TooManyRequests(''))
+    assert not retry.if_transient_error(exceptions.InvalidArgument(''))
+
+
+def test_exponential_sleep_generator_base_2():
+    gen = retry.exponential_sleep_generator(
+        1, 60, 2, jitter=0.0)
+
+    result = list(itertools.islice(gen, 8))
+    assert result == [1, 2, 4, 8, 16, 32, 60, 60]
+
+
+@mock.patch('random.uniform')
+def test_exponential_sleep_generator_jitter(uniform):
+    uniform.return_value = 1
+    gen = retry.exponential_sleep_generator(
+        1, 60, 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(
+    'google.api.core.helpers.datetime_helpers.utcnow',
+    return_value=datetime.datetime.min)
+def test_retry_target_success(utcnow, sleep):
+    predicate = retry.if_exception_type(ValueError)
+    call_count = [0]
+
+    def target():
+        call_count[0] += 1
+        if call_count[0] < 3:
+            raise ValueError()
+        return 42
+
+    result = retry.retry_target(target, predicate, range(10), None)
+
+    assert result == 42
+    assert call_count[0] == 3
+    sleep.assert_has_calls([mock.call(0), mock.call(1)])
+
+
+@mock.patch('time.sleep')
+@mock.patch(
+    'google.api.core.helpers.datetime_helpers.utcnow',
+    return_value=datetime.datetime.min)
+def test_retry_target_non_retryable_error(utcnow, sleep):
+    predicate = retry.if_exception_type(ValueError)
+    exception = TypeError()
+    target = mock.Mock(side_effect=exception)
+
+    with pytest.raises(TypeError) as exc_info:
+        retry.retry_target(target, predicate, range(10), None)
+
+    assert exc_info.value == exception
+    sleep.assert_not_called()
+
+
+@mock.patch('time.sleep')
+@mock.patch(
+    'google.api.core.helpers.datetime_helpers.utcnow')
+def test_retry_target_deadline_exceeded(utcnow, sleep):
+    predicate = retry.if_exception_type(ValueError)
+    exception = ValueError('meep')
+    target = mock.Mock(side_effect=exception)
+    # Setup the timeline so that the first call takes 5 seconds but the second
+    # call takes 6, which puts the retry over the deadline.
+    utcnow.side_effect = [
+        # The first call to utcnow establishes the start of the timeline.
+        datetime.datetime.min,
+        datetime.datetime.min + datetime.timedelta(seconds=5),
+        datetime.datetime.min + datetime.timedelta(seconds=11)]
+
+    with pytest.raises(exceptions.RetryError) as exc_info:
+        retry.retry_target(target, predicate, range(10), deadline=10)
+
+    assert exc_info.value.cause == exception
+    assert exc_info.match('Deadline of 10.0s exceeded')
+    assert exc_info.match('last exception: meep')
+    assert target.call_count == 2
+
+
+def test_retry_target_bad_sleep_generator():
+    with pytest.raises(ValueError, match='Sleep generator'):
+        retry.retry_target(
+            mock.sentinel.target, mock.sentinel.predicate, [], None)