Implement support for stripe-should-retry and retry-after headers (#621)

* Implement support for stripe-should-retry header

* Correct setup instructions in the readme

* Implement retry-after

* Some code cleanup

* Fix a lint or two

README.md

@@ -180,7 +180,7 @@ stripe-mock
 Run the following command to set up the development virtualenv:
 
 ```sh
-make init
+make
 ```
 
 Run all tests on all supported Python versions:

stripe/http_client.py

@@ -90,6 +90,7 @@ def new_default_http_client(*args, **kwargs):
 class HTTPClient(object):
     MAX_DELAY = 2
     INITIAL_DELAY = 0.5
+    MAX_RETRY_AFTER = 60
 
     def __init__(self, verify_ssl_certs=True, proxy=None):
         self._verify_ssl_certs = verify_ssl_certs
@@ -132,7 +133,7 @@ def request_with_retries(self, method, url, headers, post_data=None):
                         % connection_error.user_message
                     )
                 num_retries += 1
-                sleep_time = self._sleep_time_seconds(num_retries)
+                sleep_time = self._sleep_time_seconds(num_retries, response)
                 util.log_info(
                     (
                         "Initiating retry %i for request %s %s after "
@@ -155,24 +156,59 @@ def request(self, method, url, headers, post_data=None):
         )
 
     def _should_retry(self, response, api_connection_error, num_retries):
-        if response is not None:
-            _, status_code, _ = response
-            should_retry = status_code == 409
-        else:
+        if num_retries >= self._max_network_retries():
+            return False
+
+        if response is None:
             # We generally want to retry on timeout and connection
             # exceptions, but defer this decision to underlying subclass
             # implementations. They should evaluate the driver-specific
             # errors worthy of retries, and set flag on the error returned.
-            should_retry = api_connection_error.should_retry
-        return should_retry and num_retries < self._max_network_retries()
+            return api_connection_error.should_retry
+
+        _, status_code, rheaders = response
+
+        # The API may ask us not to retry (eg; if doing so would be a no-op)
+        # or advise us to retry (eg; in cases of lock timeouts); we defer to that.
+        #
+        # Note that we expect the headers object to be a CaseInsensitiveDict, as is the case with the requests library.
+        if rheaders is not None and "stripe-should-retry" in rheaders:
+            if rheaders["stripe-should-retry"] == "false":
+                return False
+            if rheaders["stripe-should-retry"] == "true":
+                return True
+
+        # Retry on conflict errors.
+        if status_code == 409:
+            return True
+
+        # Retry on 500, 503, and other internal errors.
+        #
+        # Note that we expect the stripe-should-retry header to be false
+        # in most cases when a 500 is returned, since our idempotency framework
+        # would typically replay it anyway.
+        if status_code >= 500:
+            return True
+
+        return False
 
     def _max_network_retries(self):
         from stripe import max_network_retries
 
         # Configured retries, isolated here for tests
         return max_network_retries
 
-    def _sleep_time_seconds(self, num_retries):
+    def _retry_after_header(self, response=None):
+        if response is None:
+            return None
+        _, _, rheaders = response
+
+        try:
+            return int(rheaders["retry-after"])
+        except (KeyError, ValueError):
+            return None
+
+    def _sleep_time_seconds(self, num_retries, response=None):
         # Apply exponential backoff with initial_network_retry_delay on the
         # number of num_retries so far as inputs.
         # Do not allow the number to exceed max_network_retry_delay.
@@ -185,6 +221,12 @@ def _sleep_time_seconds(self, num_retries):
 
         # But never sleep less than the base sleep seconds.
         sleep_seconds = max(HTTPClient.INITIAL_DELAY, sleep_seconds)
+
+        # And never sleep less than the time the API asks us to wait, assuming it's a reasonable ask.
+        retry_after = self._retry_after_header(response) or 0
+        if retry_after <= HTTPClient.MAX_RETRY_AFTER:
+            sleep_seconds = max(retry_after, sleep_seconds)
+
         return sleep_seconds
 
     def _add_jitter_time(self, sleep_seconds):

tests/test_http_client.py

@@ -92,6 +92,23 @@ def test_maximum_delay(self):
         expected = [0.5, 1.0, max_delay, max_delay, max_delay]
         self.assert_sleep_times(client, expected)
 
+    def test_retry_after_header(self):
+        client = stripe.http_client.new_default_http_client()
+        client._add_jitter_time = lambda t: t
+
+        # Prefer retry-after if it's bigger
+        assert 30 == client._sleep_time_seconds(
+            2, (None, 409, {"retry-after": "30"})
+        )
+        # Prefer default if it's bigger
+        assert 2 == client._sleep_time_seconds(
+            3, (None, 409, {"retry-after": "1"})
+        )
+        # Ignore crazy-big values
+        assert 1 == client._sleep_time_seconds(
+            2, (None, 409, {"retry-after": "300"})
+        )
+
     def test_randomness_added(self):
         client = stripe.http_client.new_default_http_client()
         random_value = 0.8
@@ -125,14 +142,20 @@ def test_should_retry_on_codes(self):
         two_xx = list(range(200, 209))
         three_xx = list(range(300, 308))
         four_xx = list(range(400, 431))
-        five_xx = list(range(500, 512))
 
         client = stripe.http_client.new_default_http_client()
-        codes = one_xx + two_xx + three_xx + four_xx + five_xx
+        client._max_network_retries = lambda: 1
+        codes = one_xx + two_xx + three_xx + four_xx
         codes.remove(409)
 
+        # These status codes should not be retried by default.
         for code in codes:
-            assert client._should_retry((None, code, None), None, 1) is False
+            assert client._should_retry((None, code, None), None, 0) is False
+
+        # These status codes should be retried by default.
+        assert client._should_retry((None, 409, None), None, 0) is True
+        assert client._should_retry((None, 500, None), None, 0) is True
+        assert client._should_retry((None, 503, None), None, 0) is True
 
     def test_should_retry_on_error(self, mocker):
         client = stripe.http_client.new_default_http_client()
@@ -145,6 +168,24 @@ def test_should_retry_on_error(self, mocker):
         api_connection_error.should_retry = False
         assert client._should_retry(None, api_connection_error, 0) is False
 
+    def test_should_retry_on_stripe_should_retry_true(self, mocker):
+        client = stripe.http_client.new_default_http_client()
+        client._max_network_retries = lambda: 1
+        headers = {"stripe-should-retry": "true"}
+
+        # Ordinarily, we would not retry a 400, but with the header as true, we would.
+        assert client._should_retry((None, 400, {}), None, 0) is False
+        assert client._should_retry((None, 400, headers), None, 0) is True
+
+    def test_should_retry_on_stripe_should_retry_false(self, mocker):
+        client = stripe.http_client.new_default_http_client()
+        client._max_network_retries = lambda: 1
+        headers = {"stripe-should-retry": "false"}
+
+        # Ordinarily, we would retry a 500, but with the header as false, we would not.
+        assert client._should_retry((None, 500, {}), None, 0) is True
+        assert client._should_retry((None, 500, headers), None, 0) is False
+
     def test_should_retry_on_num_retries(self, mocker):
         client = stripe.http_client.new_default_http_client()
         max_test_retries = 10