17
17
import abc
18
18
import random
19
19
import typing
20
+ from datetime import timedelta
20
21
21
22
from pip ._vendor .tenacity import _utils
22
23
23
24
if typing .TYPE_CHECKING :
24
25
from pip ._vendor .tenacity import RetryCallState
25
26
27
+ wait_unit_type = typing .Union [int , float , timedelta ]
28
+
29
+
30
+ def to_seconds (wait_unit : wait_unit_type ) -> float :
31
+ return float (wait_unit .total_seconds () if isinstance (wait_unit , timedelta ) else wait_unit )
32
+
26
33
27
34
class wait_base (abc .ABC ):
28
35
"""Abstract base class for wait strategies."""
@@ -44,8 +51,8 @@ def __radd__(self, other: "wait_base") -> typing.Union["wait_combine", "wait_bas
44
51
class wait_fixed (wait_base ):
45
52
"""Wait strategy that waits a fixed amount of time between each retry."""
46
53
47
- def __init__ (self , wait : float ) -> None :
48
- self .wait_fixed = wait
54
+ def __init__ (self , wait : wait_unit_type ) -> None :
55
+ self .wait_fixed = to_seconds ( wait )
49
56
50
57
def __call__ (self , retry_state : "RetryCallState" ) -> float :
51
58
return self .wait_fixed
@@ -61,9 +68,9 @@ def __init__(self) -> None:
61
68
class wait_random (wait_base ):
62
69
"""Wait strategy that waits a random amount of time between min/max."""
63
70
64
- def __init__ (self , min : typing . Union [ int , float ] = 0 , max : typing . Union [ int , float ] = 1 ) -> None : # noqa
65
- self .wait_random_min = min
66
- self .wait_random_max = max
71
+ def __init__ (self , min : wait_unit_type = 0 , max : wait_unit_type = 1 ) -> None : # noqa
72
+ self .wait_random_min = to_seconds ( min )
73
+ self .wait_random_max = to_seconds ( max )
67
74
68
75
def __call__ (self , retry_state : "RetryCallState" ) -> float :
69
76
return self .wait_random_min + (random .random () * (self .wait_random_max - self .wait_random_min ))
@@ -113,13 +120,13 @@ class wait_incrementing(wait_base):
113
120
114
121
def __init__ (
115
122
self ,
116
- start : typing . Union [ int , float ] = 0 ,
117
- increment : typing . Union [ int , float ] = 100 ,
118
- max : typing . Union [ int , float ] = _utils .MAX_WAIT , # noqa
123
+ start : wait_unit_type = 0 ,
124
+ increment : wait_unit_type = 100 ,
125
+ max : wait_unit_type = _utils .MAX_WAIT , # noqa
119
126
) -> None :
120
- self .start = start
121
- self .increment = increment
122
- self .max = max
127
+ self .start = to_seconds ( start )
128
+ self .increment = to_seconds ( increment )
129
+ self .max = to_seconds ( max )
123
130
124
131
def __call__ (self , retry_state : "RetryCallState" ) -> float :
125
132
result = self .start + (self .increment * (retry_state .attempt_number - 1 ))
@@ -142,13 +149,13 @@ class wait_exponential(wait_base):
142
149
def __init__ (
143
150
self ,
144
151
multiplier : typing .Union [int , float ] = 1 ,
145
- max : typing . Union [ int , float ] = _utils .MAX_WAIT , # noqa
152
+ max : wait_unit_type = _utils .MAX_WAIT , # noqa
146
153
exp_base : typing .Union [int , float ] = 2 ,
147
- min : typing . Union [ int , float ] = 0 , # noqa
154
+ min : wait_unit_type = 0 , # noqa
148
155
) -> None :
149
156
self .multiplier = multiplier
150
- self .min = min
151
- self .max = max
157
+ self .min = to_seconds ( min )
158
+ self .max = to_seconds ( max )
152
159
self .exp_base = exp_base
153
160
154
161
def __call__ (self , retry_state : "RetryCallState" ) -> float :
@@ -189,3 +196,37 @@ class wait_random_exponential(wait_exponential):
189
196
def __call__ (self , retry_state : "RetryCallState" ) -> float :
190
197
high = super ().__call__ (retry_state = retry_state )
191
198
return random .uniform (0 , high )
199
+
200
+
201
+ class wait_exponential_jitter (wait_base ):
202
+ """Wait strategy that applies exponential backoff and jitter.
203
+
204
+ It allows for a customized initial wait, maximum wait and jitter.
205
+
206
+ This implements the strategy described here:
207
+ https://cloud.google.com/storage/docs/retry-strategy
208
+
209
+ The wait time is min(initial * (2**n + random.uniform(0, jitter)), maximum)
210
+ where n is the retry count.
211
+ """
212
+
213
+ def __init__ (
214
+ self ,
215
+ initial : float = 1 ,
216
+ max : float = _utils .MAX_WAIT , # noqa
217
+ exp_base : float = 2 ,
218
+ jitter : float = 1 ,
219
+ ) -> None :
220
+ self .initial = initial
221
+ self .max = max
222
+ self .exp_base = exp_base
223
+ self .jitter = jitter
224
+
225
+ def __call__ (self , retry_state : "RetryCallState" ) -> float :
226
+ jitter = random .uniform (0 , self .jitter )
227
+ try :
228
+ exp = self .exp_base ** (retry_state .attempt_number - 1 )
229
+ result = self .initial * exp + jitter
230
+ except OverflowError :
231
+ result = self .max
232
+ return max (0 , min (result , self .max ))
0 commit comments