Merge pull request #17 from sot/retry

Add retry decorator and function wrapper

ska_helpers/retry/__init__.py

@@ -0,0 +1,30 @@
+"""
+Retry package initially copied from https://github.com/invl/retry.
+
+This project appears to be abandoned so moving it to ska_helpers.
+
+LICENSE:
+
+Copyright 2014 invl
+
+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.
+"""
+__all__ = ['retry', 'retry_call']
+
+import logging
+from logging import StreamHandler
+
+from .api import retry, retry_call
+
+log = logging.getLogger(__name__)
+log.addHandler(StreamHandler())

ska_helpers/retry/api.py

@@ -0,0 +1,112 @@
+import functools
+import logging
+import random
+import time
+
+
+logging_logger = logging.getLogger(__name__)
+
+
+def __retry_internal(f, exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1,
+                     jitter=0, logger=logging_logger, args=None, kwargs=None):
+    """
+    Executes a function and retries it if it failed.
+
+    :param f: the function to execute.
+    :param exceptions: an exception or a tuple of exceptions to catch. default: Exception.
+    :param tries: the maximum number of attempts. default: -1 (infinite).
+    :param delay: initial delay between attempts. default: 0.
+    :param max_delay: the maximum value of delay. default: None (no limit).
+    :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff).
+    :param jitter: extra seconds added to delay between attempts. default: 0.
+                   fixed if a number, random if a range tuple (min, max)
+    :param logger: logger.warning(fmt, error, delay) will be called on failed attempts.
+                   default: retry.logging_logger. if None, logging is disabled.
+    :param args: tuple, function args
+    :param kwargs: dict, function kwargs
+    :returns: the result of the f function.
+    """
+    _tries, _delay = tries, delay
+    while _tries:
+        try:
+            return f(*args, **kwargs)
+        except exceptions as e:
+            _tries -= 1
+            if not _tries:
+                raise
+
+            if logger is not None:
+                call_args = list(args)
+                for key, val in kwargs.items():
+                    call_args.append(f'{key}={val}')
+                call_args_str = ', '.join(str(arg) for arg in call_args)
+                func_name = getattr(f, '__name__', 'func')
+                func_call = f'{func_name}({call_args_str})'
+                logger.warning(f'WARNING: {func_call} exception: {e}, retrying '
+                               f'in {_delay} seconds...')
+
+            time.sleep(_delay)
+            _delay *= backoff
+
+            if isinstance(jitter, tuple):
+                _delay += random.uniform(*jitter)
+            else:
+                _delay += jitter
+
+            if max_delay is not None:
+                _delay = min(_delay, max_delay)
+
+
+def retry(exceptions=Exception, tries=-1, delay=0, max_delay=None, backoff=1, jitter=0,
+          logger=logging_logger):
+    """Returns a retry decorator.
+
+    :param exceptions: an exception or a tuple of exceptions to catch. default: Exception.
+    :param tries: the maximum number of attempts. default: -1 (infinite).
+    :param delay: initial delay between attempts. default: 0.
+    :param max_delay: the maximum value of delay. default: None (no limit).
+    :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff).
+    :param jitter: extra seconds added to delay between attempts. default: 0.
+                   fixed if a number, random if a range tuple (min, max)
+    :param logger: logger.warning(fmt, error, delay) will be called on failed attempts.
+                   default: retry.logging_logger. if None, logging is disabled.
+    :returns: a retry decorator.
+    """
+
+    def decorator(f):
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            return __retry_internal(f, exceptions, tries, delay, max_delay,
+                                    backoff, jitter, logger, args=args, kwargs=kwargs)
+        return wrapper
+
+    return decorator
+
+
+def retry_call(f, args=None, kwargs=None, exceptions=Exception, tries=-1, delay=0,
+               max_delay=None, backoff=1, jitter=0,
+               logger=logging_logger):
+    """
+    Calls a function and re-executes it if it failed.
+
+    :param f: the function to execute.
+    :param args: the positional arguments of the function to execute.
+    :param kwargs: the named arguments of the function to execute.
+    :param exceptions: an exception or a tuple of exceptions to catch. default: Exception.
+    :param tries: the maximum number of attempts. default: -1 (infinite).
+    :param delay: initial delay between attempts. default: 0.
+    :param max_delay: the maximum value of delay. default: None (no limit).
+    :param backoff: multiplier applied to delay between attempts. default: 1 (no backoff).
+    :param jitter: extra seconds added to delay between attempts. default: 0.
+                   fixed if a number, random if a range tuple (min, max)
+    :param logger: logger.warning(fmt, error, delay) will be called on failed attempts.
+                   default: retry.logging_logger. if None, logging is disabled.
+    :returns: the result of the f function.
+    """
+    if args is None:
+        args = []
+    if kwargs is None:
+        kwargs = {}
+
+    return __retry_internal(f, exceptions, tries, delay, max_delay,
+                            backoff, jitter, logger, args=args, kwargs=kwargs)

ska_helpers/retry/tests/test_retry.py

@@ -0,0 +1,185 @@
+try:
+    from unittest.mock import create_autospec
+except ImportError:
+    from mock import create_autospec  # noqa
+
+try:
+    from unittest.mock import MagicMock
+except ImportError:
+    from mock import MagicMock
+
+import time
+
+import pytest
+
+from ska_helpers.retry.api import retry_call
+from ska_helpers.retry import retry
+
+
+def test_retry(monkeypatch):
+    mock_sleep_time = [0]
+
+    def mock_sleep(seconds):
+        mock_sleep_time[0] += seconds
+
+    monkeypatch.setattr(time, 'sleep', mock_sleep)
+
+    hit = [0]
+
+    tries = 5
+    delay = 1
+    backoff = 2
+
+    @retry(tries=tries, delay=delay, backoff=backoff)
+    def f():
+        hit[0] += 1
+        1 / 0
+
+    with pytest.raises(ZeroDivisionError):
+        f()
+    assert hit[0] == tries
+    assert mock_sleep_time[0] == sum(
+        delay * backoff ** i for i in range(tries - 1))
+
+
+def test_tries_inf():
+    hit = [0]
+    target = 10
+
+    @retry(tries=float('inf'))
+    def f():
+        hit[0] += 1
+        if hit[0] == target:
+            return target
+        else:
+            raise ValueError
+    assert f() == target
+
+
+def test_tries_minus1():
+    hit = [0]
+    target = 10
+
+    @retry(tries=-1)
+    def f():
+        hit[0] += 1
+        if hit[0] == target:
+            return target
+        else:
+            raise ValueError
+    assert f() == target
+
+
+def test_max_delay(monkeypatch):
+    mock_sleep_time = [0]
+
+    def mock_sleep(seconds):
+        mock_sleep_time[0] += seconds
+
+    monkeypatch.setattr(time, 'sleep', mock_sleep)
+
+    hit = [0]
+
+    tries = 5
+    delay = 1
+    backoff = 2
+    max_delay = delay  # Never increase delay
+
+    @retry(tries=tries, delay=delay, max_delay=max_delay, backoff=backoff)
+    def f():
+        hit[0] += 1
+        1 / 0
+
+    with pytest.raises(ZeroDivisionError):
+        f()
+    assert hit[0] == tries
+    assert mock_sleep_time[0] == delay * (tries - 1)
+
+
+def test_fixed_jitter(monkeypatch):
+    mock_sleep_time = [0]
+
+    def mock_sleep(seconds):
+        mock_sleep_time[0] += seconds
+
+    monkeypatch.setattr(time, 'sleep', mock_sleep)
+
+    hit = [0]
+
+    tries = 10
+    jitter = 1
+
+    @retry(tries=tries, jitter=jitter)
+    def f():
+        hit[0] += 1
+        1 / 0
+
+    with pytest.raises(ZeroDivisionError):
+        f()
+    assert hit[0] == tries
+    assert mock_sleep_time[0] == sum(range(tries - 1))
+
+
+def test_retry_call():
+    f_mock = MagicMock(side_effect=RuntimeError)
+    tries = 2
+    try:
+        retry_call(f_mock, exceptions=RuntimeError, tries=tries)
+    except RuntimeError:
+        pass
+
+    assert f_mock.call_count == tries
+
+
+def test_retry_call_2():
+    side_effect = [RuntimeError, RuntimeError, 3]
+    f_mock = MagicMock(side_effect=side_effect)
+    tries = 5
+    result = None
+    try:
+        result = retry_call(f_mock, exceptions=RuntimeError, tries=tries)
+    except RuntimeError:
+        pass
+
+    assert result == 3
+    assert f_mock.call_count == len(side_effect)
+
+
+def test_retry_call_with_args():
+
+    def f(value=0):
+        if value < 0:
+            return value
+        else:
+            raise RuntimeError
+
+    return_value = -1
+    result = None
+    f_mock = MagicMock(spec=f, return_value=return_value)
+    try:
+        result = retry_call(f_mock, args=[return_value])
+    except RuntimeError:
+        pass
+
+    assert result == return_value
+    assert f_mock.call_count == 1
+
+
+def test_retry_call_with_kwargs():
+
+    def f(value=0):
+        if value < 0:
+            return value
+        else:
+            raise RuntimeError
+
+    kwargs = {'value': -1}
+    result = None
+    f_mock = MagicMock(spec=f, return_value=kwargs['value'])
+    try:
+        result = retry_call(f_mock, kwargs=kwargs)
+    except RuntimeError:
+        pass
+
+    assert result == kwargs['value']
+    assert f_mock.call_count == 1