From 84763d2ff434ad21e6c96bddcbbe15e6c7fd9610 Mon Sep 17 00:00:00 2001
From: David Leong <116610336+leongdl@users.noreply.github.com>
Date: Tue, 24 Sep 2024 17:17:48 -0700
Subject: [PATCH] feat: Add a simple backoff and retry utility helper

Signed-off-by: David Leong <116610336+leongdl@users.noreply.github.com>
---
 src/deadline/job_attachments/_utils.py        | 58 ++++++++++++++++++-
 .../deadline_job_attachments/test_utils.py    | 21 +++++++
 2 files changed, 78 insertions(+), 1 deletion(-)

diff --git a/src/deadline/job_attachments/_utils.py b/src/deadline/job_attachments/_utils.py
index 5a148a810..b9694eb76 100644
--- a/src/deadline/job_attachments/_utils.py
+++ b/src/deadline/job_attachments/_utils.py
@@ -1,9 +1,12 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 
 import datetime
+from functools import wraps
 from hashlib import shake_256
 from pathlib import Path
-from typing import Tuple, Union
+import random
+import time
+from typing import Any, Callable, Optional, Tuple, Type, Union
 import uuid
 import ctypes
 import sys
@@ -98,3 +101,56 @@ def _is_windows_file_path_limit() -> bool:
 
         return bool(ntdll.RtlAreLongPathsEnabled())
     return True
+
+
+def _retry(
+    ExceptionToCheck: Union[Type[Exception], Tuple[Type[Exception], ...]] = AssertionError,
+    tries: int = 4,
+    delay: Union[int, float, Tuple[Union[int, float], Union[int, float]]] = 3.0,
+    backoff: float = 2.0,
+    logger: Optional[Callable] = print,
+) -> Callable:
+    """Retry calling the decorated function using an exponential backoff.
+
+    http://www.saltycrane.com/blog/2009/11/trying-out-retry-decorator-python/
+    original from: http://wiki.python.org/moin/PythonDecoratorLibrary#Retry
+
+    :param ExceptionToCheck: the exception to check. may be a tuple of
+        exceptions to check
+    :type ExceptionToCheck: Exception or tuple
+    :param tries: number of times to try (not retry) before giving up
+    :type tries: int
+    :param delay: initial delay between retries in seconds
+    :type delay: float or tuple
+    :param backoff: backoff multiplier e.g. value of 2 will double the delay
+        each retry
+    :type backoff: float
+    :param logger: logging function to use. If None, won't log
+    :type logger: logging.Logger instance
+    """
+
+    def deco_retry(f: Callable) -> Callable:
+        @wraps(f)
+        def f_retry(*args: Any, **kwargs: Any) -> Callable:
+            mtries: int = tries
+            if isinstance(delay, (float, int)):
+                mdelay = delay
+            elif isinstance(delay, tuple):
+                mdelay = random.uniform(delay[0], delay[1])
+            else:
+                raise ValueError(f"Provided delay {delay} isn't supported")
+
+            while mtries > 1:
+                try:
+                    return f(*args, **kwargs)
+                except ExceptionToCheck as e:
+                    if logger:
+                        logger(f"{str(e)}, Retrying in {mdelay} seconds...")
+                    time.sleep(mdelay)
+                    mtries -= 1
+                    mdelay *= backoff
+            return f(*args, **kwargs)
+
+        return f_retry  # true decorator
+
+    return deco_retry
diff --git a/test/unit/deadline_job_attachments/test_utils.py b/test/unit/deadline_job_attachments/test_utils.py
index b33fb3389..808b29ef9 100644
--- a/test/unit/deadline_job_attachments/test_utils.py
+++ b/test/unit/deadline_job_attachments/test_utils.py
@@ -7,6 +7,7 @@
 
 from deadline.job_attachments._utils import (
     _is_relative_to,
+    _retry,
 )
 
 
@@ -60,3 +61,23 @@ def test_is_relative_to_on_windows(self, path1, path2, expected):
         Tests if the is_relative_to() works correctly when using Windows paths.
         """
         assert _is_relative_to(path1, path2) == expected
+
+    def test_retry(self):
+        """
+        Test a function that throws an exception is retried.
+        """
+        call_count = 0
+
+        # Given
+        @_retry(ExceptionToCheck=NotImplementedError, tries=2, delay=0.1, backoff=0.1)
+        def test_bad_function():
+            nonlocal call_count
+            call_count = call_count + 1
+            if call_count == 1:
+                raise NotImplementedError()
+
+        # When
+        test_bad_function()
+
+        # Then
+        assert call_count == 2