1# Copyright 2022 Google LLC
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import asyncio
16import random
17import time
18
19from google.auth import exceptions
20
21# The default amount of retry attempts
22_DEFAULT_RETRY_TOTAL_ATTEMPTS = 3
23
24# The default initial backoff period (1.0 second).
25_DEFAULT_INITIAL_INTERVAL_SECONDS = 1.0
26
27# The default randomization factor (0.1 which results in a random period ranging
28# between 10% below and 10% above the retry interval).
29_DEFAULT_RANDOMIZATION_FACTOR = 0.1
30
31# The default multiplier value (2 which is 100% increase per back off).
32_DEFAULT_MULTIPLIER = 2.0
33
34"""Exponential Backoff Utility
35
36This is a private module that implements the exponential back off algorithm.
37It can be used as a utility for code that needs to retry on failure, for example
38an HTTP request.
39"""
40
41
42class _BaseExponentialBackoff:
43 """An exponential backoff iterator base class.
44
45 Args:
46 total_attempts Optional[int]:
47 The maximum amount of retries that should happen.
48 The default value is 3 attempts.
49 initial_wait_seconds Optional[int]:
50 The amount of time to sleep in the first backoff. This parameter
51 should be in seconds.
52 The default value is 1 second.
53 randomization_factor Optional[float]:
54 The amount of jitter that should be in each backoff. For example,
55 a value of 0.1 will introduce a jitter range of 10% to the
56 current backoff period.
57 The default value is 0.1.
58 multiplier Optional[float]:
59 The backoff multipler. This adjusts how much each backoff will
60 increase. For example a value of 2.0 leads to a 200% backoff
61 on each attempt. If the initial_wait is 1.0 it would look like
62 this sequence [1.0, 2.0, 4.0, 8.0].
63 The default value is 2.0.
64 """
65
66 def __init__(
67 self,
68 total_attempts=_DEFAULT_RETRY_TOTAL_ATTEMPTS,
69 initial_wait_seconds=_DEFAULT_INITIAL_INTERVAL_SECONDS,
70 randomization_factor=_DEFAULT_RANDOMIZATION_FACTOR,
71 multiplier=_DEFAULT_MULTIPLIER,
72 ):
73 if total_attempts < 1:
74 raise exceptions.InvalidValue(
75 f"total_attempts must be greater than or equal to 1 but was {total_attempts}"
76 )
77
78 self._total_attempts = total_attempts
79 self._initial_wait_seconds = initial_wait_seconds
80
81 self._current_wait_in_seconds = self._initial_wait_seconds
82
83 self._randomization_factor = randomization_factor
84 self._multiplier = multiplier
85 self._backoff_count = 0
86
87 @property
88 def total_attempts(self):
89 """The total amount of backoff attempts that will be made."""
90 return self._total_attempts
91
92 @property
93 def backoff_count(self):
94 """The current amount of backoff attempts that have been made."""
95 return self._backoff_count
96
97 def _reset(self):
98 self._backoff_count = 0
99 self._current_wait_in_seconds = self._initial_wait_seconds
100
101 def _calculate_jitter(self):
102 jitter_variance = self._current_wait_in_seconds * self._randomization_factor
103 jitter = random.uniform(
104 self._current_wait_in_seconds - jitter_variance,
105 self._current_wait_in_seconds + jitter_variance,
106 )
107
108 return jitter
109
110
111class ExponentialBackoff(_BaseExponentialBackoff):
112 """An exponential backoff iterator. This can be used in a for loop to
113 perform requests with exponential backoff.
114 """
115
116 def __init__(self, *args, **kwargs):
117 super(ExponentialBackoff, self).__init__(*args, **kwargs)
118
119 def __iter__(self):
120 self._reset()
121 return self
122
123 def __next__(self):
124 if self._backoff_count >= self._total_attempts:
125 raise StopIteration
126 self._backoff_count += 1
127
128 if self._backoff_count <= 1:
129 return self._backoff_count
130
131 jitter = self._calculate_jitter()
132
133 time.sleep(jitter)
134
135 self._current_wait_in_seconds *= self._multiplier
136 return self._backoff_count
137
138
139class AsyncExponentialBackoff(_BaseExponentialBackoff):
140 """An async exponential backoff iterator. This can be used in a for loop to
141 perform async requests with exponential backoff.
142 """
143
144 def __init__(self, *args, **kwargs):
145 super(AsyncExponentialBackoff, self).__init__(*args, **kwargs)
146
147 def __aiter__(self):
148 self._reset()
149 return self
150
151 async def __anext__(self):
152 if self._backoff_count >= self._total_attempts:
153 raise StopAsyncIteration
154 self._backoff_count += 1
155
156 if self._backoff_count <= 1:
157 return self._backoff_count
158
159 jitter = self._calculate_jitter()
160
161 await asyncio.sleep(jitter)
162
163 self._current_wait_in_seconds *= self._multiplier
164 return self._backoff_count