feat: Implement rate-limit retries in RestClient

evansims · evansims · commit 7f810a72ba3f · 2021-08-16T18:02:36.000-04:00
diff --git a/auth0/v3/management/rest.py b/auth0/v3/management/rest.py
@@ -2,14 +2,14 @@
 import json
 import platform
 import sys
-
 import requests
 
 from ..exceptions import Auth0Error, RateLimitError
+from time import sleep
+from random import randint
 
 UNKNOWN_ERROR = 'a0.sdk.internal.unknown'
 
-
 class RestClient(object):
     """Provides simple methods for handling all RESTful api endpoints.
 
@@ -22,14 +22,18 @@ class RestClient(object):
             (defaults to 5.0 for both)
     """
 
-    def __init__(self, jwt, telemetry=True, timeout=5.0):
+    def __init__(self, jwt, telemetry=True, timeout=5.0, retries=3):
         self.jwt = jwt
         self.timeout = timeout
+        self.retries = retries
+        self._metrics = {'retries': 0, 'backoff': []}
+        self._skip_sleep = False
 
         self.base_headers = {
             'Authorization': 'Bearer {}'.format(self.jwt),
             'Content-Type': 'application/json',
         }
+
         if telemetry:
             py_version = platform.python_version()
             version = sys.modules['auth0'].__version__
@@ -47,10 +51,69 @@ def __init__(self, jwt, telemetry=True, timeout=5.0):
                 'Auth0-Client': base64.b64encode(auth0_client),
             })
 
+    # Returns a hard cap for the maximum number of retries allowed (10)
+    def MAX_REQUEST_RETRIES(self):
+        return 10
+
+    # Returns the maximum amount of jitter to introduce in milliseconds (100ms)
+    def MAX_REQUEST_RETRY_JITTER(self):
+        return 100
+
+    # Returns the maximum delay window allowed (1000ms)
+    def MAX_REQUEST_RETRY_DELAY(self):
+        return 1000
+
+    # Returns the minimum delay window allowed (100ms)
+    def MIN_REQUEST_RETRY_DELAY(self):
+        return 100
+
     def get(self, url, params=None):
         headers = self.base_headers.copy()
 
-        response = requests.get(url, params=params, headers=headers, timeout=self.timeout)
+        # Track the API request attempt number
+        attempt = 0
+
+        # Reset the metrics tracker
+        self._metrics = {'retries': 0, 'backoff': []}
+
+        # Cap the maximum number of retries to 10 or fewer. Floor the retries at 0.
+        retries = min(self.MAX_REQUEST_RETRIES(), max(0, self.retries))
+
+        while True:
+            # Increment attempt number
+            attempt += 1
+
+            # Issue the request
+            response = requests.get(url, params=params, headers=headers, timeout=self.timeout);
+
+            # If the response did not have a 429 header, or the retries were configured at 0, or the attempt number is equal to or greater than the configured retries, break
+            if response.status_code != 429 or retries <= 0 or attempt > retries:
+                break
+
+            # Retry the request. Apply a exponential backoff for subsequent attempts, using this formula:
+            # max(MIN_REQUEST_RETRY_DELAY, min(MAX_REQUEST_RETRY_DELAY, (100ms * (2 ** attempt - 1)) + random_between(1, MAX_REQUEST_RETRY_JITTER)))
+
+            # ✔ Increases base delay by (100ms * (2 ** attempt - 1))
+            wait = 100 * 2 ** (attempt - 1)
+
+            # ✔ Introduces jitter to the base delay; increases delay between 1ms to MAX_REQUEST_RETRY_JITTER (100ms)
+            wait += randint(1, self.MAX_REQUEST_RETRY_JITTER())
+
+            # ✔ Is never more than MAX_REQUEST_RETRY_DELAY (1s)
+            wait = min(self.MAX_REQUEST_RETRY_DELAY(), wait)
+
+            # ✔ Is never less than MIN_REQUEST_RETRY_DELAY (100ms)
+            wait = max(self.MIN_REQUEST_RETRY_DELAY(), wait)
+
+            self._metrics['retries'] = attempt
+            self._metrics['backoff'].append(wait)
+
+            # Skip calling sleep() when running unit tests
+            if self._skip_sleep is False:
+                # sleep() functions in seconds, so convert the milliseconds formula above accordingly
+                sleep(wait / 1000)
+
+        # Return the final Response
         return self._process_response(response)
 
     def post(self, url, data=None):
