From da1bfc9cfdf259b5f28c0f89408065b03afb9894 Mon Sep 17 00:00:00 2001
From: Isaac Good <IsaacG@users.noreply.github.com>
Date: Wed, 27 Apr 2022 09:43:10 -0700
Subject: [PATCH] Implement a wait.wait_exponential_jitter per Google's storage
 guide (#351)

* Implement a wait.wait_exponential_jitter per Google's storage retry guide

* Define a ClientError so Sphinx does not fail

* Fix spelling typos

* Simplify typing, replacing `int | float` with `float`

* Drop needless `#noqa`

Co-authored-by: Isaac Good <goodi@twosigma.com>
---
 ...t_exponential_jitter-6ffc81dddcbaa6d3.yaml |  5 +++
 tenacity/__init__.py                          |  1 +
 tenacity/wait.py                              | 34 +++++++++++++++++++
 tests/test_tenacity.py                        | 22 ++++++++++++
 4 files changed, 62 insertions(+)
 create mode 100644 releasenotes/notes/wait_exponential_jitter-6ffc81dddcbaa6d3.yaml

diff --git a/releasenotes/notes/wait_exponential_jitter-6ffc81dddcbaa6d3.yaml b/releasenotes/notes/wait_exponential_jitter-6ffc81dddcbaa6d3.yaml
new file mode 100644
index 00000000..870380ca
--- /dev/null
+++ b/releasenotes/notes/wait_exponential_jitter-6ffc81dddcbaa6d3.yaml
@@ -0,0 +1,5 @@
+---
+features:
+  - |
+    Implement a wait.wait_exponential_jitter per Google's storage retry guide.
+    See https://cloud.google.com/storage/docs/retry-strategy
diff --git a/tenacity/__init__.py b/tenacity/__init__.py
index 1e3ff865..fd403761 100644
--- a/tenacity/__init__.py
+++ b/tenacity/__init__.py
@@ -63,6 +63,7 @@
 from .wait import wait_random  # noqa
 from .wait import wait_random_exponential  # noqa
 from .wait import wait_random_exponential as wait_full_jitter  # noqa
+from .wait import wait_exponential_jitter  # noqa
 
 # Import all built-in before strategies for easier usage.
 from .before import before_log  # noqa
diff --git a/tenacity/wait.py b/tenacity/wait.py
index aacb58d6..289705c7 100644
--- a/tenacity/wait.py
+++ b/tenacity/wait.py
@@ -189,3 +189,37 @@ class wait_random_exponential(wait_exponential):
     def __call__(self, retry_state: "RetryCallState") -> float:
         high = super().__call__(retry_state=retry_state)
         return random.uniform(0, high)
+
+
+class wait_exponential_jitter(wait_base):
+    """Wait strategy that applies exponential backoff and jitter.
+
+    It allows for a customized initial wait, maximum wait and jitter.
+
+    This implements the strategy described here:
+    https://cloud.google.com/storage/docs/retry-strategy
+
+    The wait time is min(initial * (2**n + random.uniform(0, jitter)), maximum)
+    where n is the retry count.
+    """
+
+    def __init__(
+        self,
+        initial: float = 1,
+        max: float = _utils.MAX_WAIT,  # noqa
+        exp_base: float = 2,
+        jitter: float = 1,
+    ) -> None:
+        self.initial = initial
+        self.max = max
+        self.exp_base = exp_base
+        self.jitter = jitter
+
+    def __call__(self, retry_state: "RetryCallState") -> float:
+        jitter = random.uniform(0, self.jitter)
+        try:
+            exp = self.exp_base ** (retry_state.attempt_number - 1)
+            result = self.initial * exp + jitter
+        except OverflowError:
+            result = self.max
+        return max(0, min(result, self.max))
diff --git a/tests/test_tenacity.py b/tests/test_tenacity.py
index b9016e71..d9a48583 100644
--- a/tests/test_tenacity.py
+++ b/tests/test_tenacity.py
@@ -435,6 +435,28 @@ def mean(lst):
         self._assert_inclusive_epsilon(mean(attempt[8]), 30, 2.56)
         self._assert_inclusive_epsilon(mean(attempt[9]), 30, 2.56)
 
+    def test_wait_exponential_jitter(self):
+        fn = tenacity.wait_exponential_jitter(max=60)
+
+        for _ in range(1000):
+            self._assert_inclusive_range(fn(make_retry_state(1, 0)), 1, 2)
+            self._assert_inclusive_range(fn(make_retry_state(2, 0)), 2, 3)
+            self._assert_inclusive_range(fn(make_retry_state(3, 0)), 4, 5)
+            self._assert_inclusive_range(fn(make_retry_state(4, 0)), 8, 9)
+            self._assert_inclusive_range(fn(make_retry_state(5, 0)), 16, 17)
+            self._assert_inclusive_range(fn(make_retry_state(6, 0)), 32, 33)
+            self.assertEqual(fn(make_retry_state(7, 0)), 60)
+            self.assertEqual(fn(make_retry_state(8, 0)), 60)
+            self.assertEqual(fn(make_retry_state(9, 0)), 60)
+
+        fn = tenacity.wait_exponential_jitter(10, 5)
+        for _ in range(1000):
+            self.assertEqual(fn(make_retry_state(1, 0)), 5)
+
+        # Default arguments exist
+        fn = tenacity.wait_exponential_jitter()
+        fn(make_retry_state(0, 0))
+
     def test_wait_retry_state_attributes(self):
         class ExtractCallState(Exception):
             pass