From 74aa7e20fba61ea17ea391b140229c736a770a7e Mon Sep 17 00:00:00 2001
From: Jason Lubken <jlubken@users.noreply.github.com>
Date: Wed, 22 Jan 2020 11:36:24 -0500
Subject: [PATCH 1/3] Add retry decorator

---
 src/dsdk/utils.py | 51 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/src/dsdk/utils.py b/src/dsdk/utils.py
index a9d68af..8691be9 100644
--- a/src/dsdk/utils.py
+++ b/src/dsdk/utils.py
@@ -6,6 +6,10 @@
 import pickle
 from collections import OrderedDict
 from datetime import datetime
+from functools import wraps
+from logging import NullHandler, getLogger
+from time import sleep as default_sleep
+from typing import Callable, Sequence
 from warnings import warn
 
 from configargparse import ArgParser
@@ -29,6 +33,10 @@
     MongoClient = None
 
 
+logger = getLogger(__name__)
+logger.addHandler(NullHandler())
+
+
 def get_base_config() -> ArgParser:
     """Get the base configuration parser."""
     config_parser = ArgParser(
@@ -124,3 +132,46 @@ def __setitem__(self, key, value):
         if key in self:
             raise KeyError("{} has already been set".format(key))
         super(WriteOnceDict, self).__setitem__(key, value)
+
+
+def retry(
+    exceptions: Sequence[Exception],
+    retries: int = 5,
+    delay: int = 1,
+    backoff: float = 1.5,
+    sleep: Callable = default_sleep,
+):
+    """
+    Retry calling the decorated function using an exponential backoff.
+
+    Args:
+        exceptions: The exception to check. may be a tuple of
+            exceptions to check.
+        retries: Number of times to retry before giving up.
+        delay: Initial delay between retries in seconds.
+        backoff: Backoff multiplier (e.g. value of 2 will double the delay
+            each retry).
+    """
+
+    def wrapper(func):
+        @wraps(func)
+        def wrapped(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except exceptions as exception:
+                logger.exception(exception)
+                wait = delay
+                for _ in range(retries):
+                    message = f"Retrying in {wait} seconds..."
+                    logger.warning(message)
+                    sleep(wait)
+                    wait *= backoff
+                    try:
+                        return func(*args, **kwargs)
+                    except exceptions as exception:
+                        logger.exception(exception)
+                raise
+
+        return wrapped
+
+    return wrapper

From e3d385f29a6b79488ae0dff8307f3b4b4901f4d6 Mon Sep 17 00:00:00 2001
From: Jason Lubken <jlubken@users.noreply.github.com>
Date: Wed, 22 Jan 2020 13:54:27 -0500
Subject: [PATCH 2/3] Add retry tests

---
 src/dsdk/utils.py  |  4 ++-
 tests/test_dsdk.py | 65 ++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

diff --git a/src/dsdk/utils.py b/src/dsdk/utils.py
index 8691be9..ec68058 100644
--- a/src/dsdk/utils.py
+++ b/src/dsdk/utils.py
@@ -137,7 +137,7 @@ def __setitem__(self, key, value):
 def retry(
     exceptions: Sequence[Exception],
     retries: int = 5,
-    delay: int = 1,
+    delay: float = 1.0,
     backoff: float = 1.5,
     sleep: Callable = default_sleep,
 ):
@@ -152,6 +152,8 @@ def retry(
         backoff: Backoff multiplier (e.g. value of 2 will double the delay
             each retry).
     """
+    delay = float(delay)
+    backoff = float(backoff)
 
     def wrapper(func):
         @wraps(func)
diff --git a/tests/test_dsdk.py b/tests/test_dsdk.py
index a8399e2..c3b79a6 100644
--- a/tests/test_dsdk.py
+++ b/tests/test_dsdk.py
@@ -6,6 +6,7 @@
 import configargparse
 
 from dsdk import BaseBatchJob, Block
+from dsdk.utils import retry
 
 
 def test_batch(monkeypatch):
@@ -23,3 +24,67 @@ def run(self):
     batch.run()
     assert len(batch.evidence) == 1
     assert batch.evidence["test"] == 42
+
+
+def test_retry_other_exception():
+    """Test retry other exception."""
+
+    exceptions_in = [
+        RuntimeError("what?"),
+        NotImplementedError("how?"),
+        RuntimeError("no!"),
+    ]
+    actual = []
+    expected = [1.0, 1.5, 2.25]
+
+    def sleep(wait: float):
+        actual.append(wait)
+
+    @retry(
+        (NotImplementedError, RuntimeError),
+        retries=4,
+        delay=1.0,
+        backoff=1.5,
+        sleep=sleep,
+    )
+    def explode():
+        raise exceptions_in.pop()
+
+    try:
+        explode()
+        raise AssertionError("IndexError expected")
+    except IndexError:
+        assert actual == expected
+
+
+def test_retry_exhausted():
+    """Test retry."""
+
+    exceptions_in = [
+        RuntimeError("what?"),
+        NotImplementedError("how?"),
+        RuntimeError("no!"),
+        NotImplementedError("when?"),
+    ]
+    actual = []
+    expected = [1.0, 1.5]
+
+    def sleep(wait: float):
+        actual.append(wait)
+
+    @retry(
+        (NotImplementedError, RuntimeError),
+        retries=2,
+        delay=1.0,
+        backoff=1.5,
+        sleep=sleep,
+    )
+    def explode():
+        raise exceptions_in.pop()
+
+    try:
+        explode()
+        raise AssertionError("NotImplementedError expected")
+    except NotImplementedError as exception:
+        assert actual == expected
+        assert str(exception) == "when?"

From 50c0a2e00a3f64358beed72eab6821fabd3c6269 Mon Sep 17 00:00:00 2001
From: Jason Lubken <jlubken@users.noreply.github.com>
Date: Wed, 22 Jan 2020 14:08:29 -0500
Subject: [PATCH 3/3] Format retry wait time

---
 src/dsdk/utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/dsdk/utils.py b/src/dsdk/utils.py
index ec68058..9debc0d 100644
--- a/src/dsdk/utils.py
+++ b/src/dsdk/utils.py
@@ -164,7 +164,7 @@ def wrapped(*args, **kwargs):
                 logger.exception(exception)
                 wait = delay
                 for _ in range(retries):
-                    message = f"Retrying in {wait} seconds..."
+                    message = f"Retrying in {wait:.2f} seconds..."
                     logger.warning(message)
                     sleep(wait)
                     wait *= backoff