feat: allow RetryStrategy to be configured with a custom initial wait…

… and multiplier (#216)

* feat: allow RetryStrategy to be configured with a custom initial wait and multiplier

* lint

* add coverage for initial delay and multiplier args

* comments for tests

google/resumable_media/_helpers.py

@@ -106,7 +106,7 @@ def require_status_code(response, status_codes, get_status_code, callback=do_not
     return status_code
 
 
-def calculate_retry_wait(base_wait, max_sleep):
+def calculate_retry_wait(base_wait, max_sleep, multiplier=2.0):
     """Calculate the amount of time to wait before a retry attempt.
 
     Wait time grows exponentially with the number of attempts, until
@@ -117,15 +117,16 @@ def calculate_retry_wait(base_wait, max_sleep):
 
     Args:
         base_wait (float): The "base" wait time (i.e. without any jitter)
-            that will be doubled until it reaches the maximum sleep.
+            that will be multiplied until it reaches the maximum sleep.
         max_sleep (float): Maximum value that a sleep time is allowed to be.
+        multiplier (float): Multiplier to apply to the base wait.
 
     Returns:
         Tuple[float, float]: The new base wait time as well as the wait time
         to be applied (with a random amount of jitter between 0 and 1 seconds
         added).
     """
-    new_base_wait = 2.0 * base_wait
+    new_base_wait = multiplier * base_wait
     if new_base_wait > max_sleep:
         new_base_wait = max_sleep
 
@@ -157,7 +158,8 @@ def wait_and_retry(func, get_status_code, retry_strategy):
     """
     total_sleep = 0.0
     num_retries = 0
-    base_wait = 0.5  # When doubled will give 1.0
+    # base_wait will be multiplied by the multiplier on the first retry.
+    base_wait = float(retry_strategy.initial_delay) / retry_strategy.multiplier
 
     # Set the retriable_exception_type if possible. We expect requests to be
     # present here and the transport to be using requests.exceptions errors,
@@ -187,7 +189,9 @@ def wait_and_retry(func, get_status_code, retry_strategy):
 
             return response
 
-        base_wait, wait_time = calculate_retry_wait(base_wait, retry_strategy.max_sleep)
+        base_wait, wait_time = calculate_retry_wait(
+            base_wait, retry_strategy.max_sleep, retry_strategy.multiplier
+        )
 
         num_retries += 1
         total_sleep += wait_time

google/resumable_media/common.py

@@ -119,20 +119,29 @@ class RetryStrategy(object):
         max_cumulative_retry (Optional[float]): The maximum **total** amount of
             time to sleep during retry process.
         max_retries (Optional[int]): The number of retries to attempt.
+        initial_delay (Optional[float]): The initial delay. Default 1.0 second.
+        muiltiplier (Optional[float]): Exponent of the backoff. Default is 2.0.
 
     Attributes:
         max_sleep (float): Maximum amount of time allowed between requests.
         max_cumulative_retry (Optional[float]): Maximum total sleep time
             allowed during retry process.
         max_retries (Optional[int]): The number retries to attempt.
+        initial_delay (Optional[float]): The initial delay. Default 1.0 second.
+        muiltiplier (Optional[float]): Exponent of the backoff. Default is 2.0.
 
     Raises:
         ValueError: If both of ``max_cumulative_retry`` and ``max_retries``
             are passed.
     """
 
     def __init__(
-        self, max_sleep=MAX_SLEEP, max_cumulative_retry=None, max_retries=None
+        self,
+        max_sleep=MAX_SLEEP,
+        max_cumulative_retry=None,
+        max_retries=None,
+        initial_delay=1.0,
+        multiplier=2.0,
     ):
         if max_cumulative_retry is not None and max_retries is not None:
             raise ValueError(_SLEEP_RETRY_ERROR_MSG)
@@ -142,6 +151,8 @@ def __init__(
         self.max_sleep = max_sleep
         self.max_cumulative_retry = max_cumulative_retry
         self.max_retries = max_retries
+        self.initial_delay = initial_delay
+        self.multiplier = multiplier
 
     def retry_allowed(self, total_sleep, num_retries):
         """Check if another retry is allowed.

tests/unit/test__helpers.py

@@ -150,6 +150,14 @@ def test_under_limit(self, randint_mock):
         assert wait_time == 32.875
         randint_mock.assert_called_once_with(0, 1000)
 
+    @mock.patch(u"random.randint", return_value=875)
+    def test_custom_multiplier(self, randint_mock):
+        base_wait, wait_time = _helpers.calculate_retry_wait(16.0, 64.0, 3)
+
+        assert base_wait == 48.0
+        assert wait_time == 48.875
+        randint_mock.assert_called_once_with(0, 1000)
+
 
 class Test_wait_and_retry(object):
     def test_success_no_retry(self):
@@ -195,6 +203,41 @@ def test_success_with_retry(self, randint_mock, sleep_mock):
         sleep_mock.assert_any_call(2.625)
         sleep_mock.assert_any_call(4.375)
 
+    @mock.patch(u"time.sleep")
+    @mock.patch(u"random.randint")
+    def test_success_with_retry_custom_delay(self, randint_mock, sleep_mock):
+        randint_mock.side_effect = [125, 625, 375]
+
+        status_codes = (
+            http_client.INTERNAL_SERVER_ERROR,
+            http_client.BAD_GATEWAY,
+            http_client.SERVICE_UNAVAILABLE,
+            http_client.NOT_FOUND,
+        )
+        responses = [_make_response(status_code) for status_code in status_codes]
+        func = mock.Mock(side_effect=responses, spec=[])
+
+        retry_strategy = common.RetryStrategy(initial_delay=3.0, multiplier=4)
+        ret_val = _helpers.wait_and_retry(func, _get_status_code, retry_strategy)
+
+        assert ret_val == responses[-1]
+        assert status_codes[-1] not in common.RETRYABLE
+
+        assert func.call_count == 4
+        assert func.mock_calls == [mock.call()] * 4
+
+        assert randint_mock.call_count == 3
+        assert randint_mock.mock_calls == [mock.call(0, 1000)] * 3
+
+        assert sleep_mock.call_count == 3
+        sleep_mock.assert_any_call(3.125)  # initial delay 3 + jitter 0.125
+        sleep_mock.assert_any_call(
+            12.625
+        )  # previous delay 3 * multiplier 4 + jitter 0.625
+        sleep_mock.assert_any_call(
+            48.375
+        )  # previous delay 12 * multiplier 4 + jitter 0.375
+
     @mock.patch(u"time.sleep")
     @mock.patch(u"random.randint")
     def test_success_with_retry_connection_error(self, randint_mock, sleep_mock):

tests/unit/test_common.py

@@ -40,6 +40,14 @@ def test_constructor_failure(self):
 
         exc_info.match(common._SLEEP_RETRY_ERROR_MSG)
 
+    def test_constructor_custom_delay_and_multiplier(self):
+        retry_strategy = common.RetryStrategy(initial_delay=3.0, multiplier=4)
+        assert retry_strategy.max_sleep == common.MAX_SLEEP
+        assert retry_strategy.max_cumulative_retry == common.MAX_CUMULATIVE_RETRY
+        assert retry_strategy.max_retries is None
+        assert retry_strategy.initial_delay == 3.0
+        assert retry_strategy.multiplier == 4
+
     def test_constructor_explicit_bound_cumulative(self):
         max_sleep = 10.0
         max_cumulative_retry = 100.0