Support server sending back Retry-After header for rate limiting

raven/base.py

@@ -61,27 +61,30 @@ def __init__(self):
         self.status = self.ONLINE
         self.last_check = None
         self.retry_number = 0
+        self.retry_after = 0
 
     def should_try(self):
         if self.status == self.ONLINE:
             return True
 
-        interval = min(self.retry_number, 6) ** 2
+        interval = self.retry_after or min(self.retry_number, 6) ** 2
 
         if time.time() - self.last_check > interval:
             return True
 
         return False
 
-    def set_fail(self):
+    def set_fail(self, retry_after=0):
         self.status = self.ERROR
         self.retry_number += 1
         self.last_check = time.time()
+        self.retry_after = retry_after
 
     def set_success(self):
         self.status = self.ONLINE
         self.last_check = None
         self.retry_number = 0
+        self.retry_after = 0
 
     def did_fail(self):
         return self.status == self.ERROR
@@ -522,10 +525,15 @@ def _successful_send(self):
         self.state.set_success()
 
     def _failed_send(self, e, url, data):
+        retry_after = 0
         if isinstance(e, APIError):
             self.error_logger.error('Unable to capture event: %s', e.message)
         elif isinstance(e, HTTPError):
             body = e.read()
+            try:
+                retry_after = int(e.headers.get('Retry-After'))
+            except (ValueError, TypeError):
+                pass
             self.error_logger.error(
                 'Unable to reach Sentry log server: %s (url: %s, body: %s)',
                 e, url, body, exc_info=True,
@@ -537,7 +545,7 @@ def _failed_send(self, e, url, data):
 
         message = self._get_log_message(data)
         self.error_logger.error('Failed to submit message: %r', message)
-        self.state.set_fail()
+        self.state.set_fail(retry_after=retry_after)
 
     def send_remote(self, url, data, headers=None):
         if headers is None:

tests/base/tests.py

@@ -6,12 +6,14 @@
 import pytest
 import raven
 import time
+from StringIO import StringIO
 from socket import socket, AF_INET, SOCK_DGRAM
 from raven.base import Client, ClientState
 from raven.transport import AsyncTransport
 from raven.utils.stacks import iter_stack_frames
 from raven.utils import six
 from raven.utils.testutils import TestCase
+from raven.utils.compat import HTTPError
 
 
 class TempStoreClient(Client):
@@ -62,6 +64,22 @@ def test_set_success(self):
         self.assertEquals(state.last_check, None)
         self.assertEquals(state.retry_number, 0)
 
+    def test_should_try_retry_after(self):
+        state = ClientState()
+        state.status = state.ERROR
+        state.last_check = time.time()
+        state.retry_number = 1
+        state.retry_after = 1
+        self.assertFalse(state.should_try())
+
+    def test_should_try_retry_after_passed(self):
+        state = ClientState()
+        state.status = state.ERROR
+        state.last_check = time.time() - 1
+        state.retry_number = 1
+        state.retry_after = 1
+        self.assertTrue(state.should_try())
+
 
 class ClientTest(TestCase):
     def setUp(self):
@@ -96,6 +114,31 @@ def test_send_remote_failover(self, should_try, send):
         client.send_remote('sync+http://example.com/api/store', 'foo')
         self.assertEquals(client.state.status, client.state.ONLINE)
 
+    @mock.patch('raven.transport.http.HTTPTransport.send')
+    @mock.patch('raven.base.ClientState.should_try')
+    def test_send_remote_failover_with_retry_after(self, should_try, send):
+        should_try.return_value = True
+
+        client = Client(
+            dsn='sync+http://public:[email protected]/1'
+        )
+
+        e = HTTPError(
+            'http://example.com/api/store', 429, 'oops',
+            {'Retry-After': '5'}, StringIO())
+
+        # test error
+        send.side_effect = e
+        client.send_remote('sync+http://example.com/api/store', 'foo')
+        self.assertEquals(client.state.status, client.state.ERROR)
+        self.assertEqual(client.state.retry_after, 5)
+
+        # test recovery
+        send.side_effect = None
+        client.send_remote('sync+http://example.com/api/store', 'foo')
+        self.assertEquals(client.state.status, client.state.ONLINE)
+        self.assertEqual(client.state.retry_after, 0)
+
     @mock.patch('raven.base.Client._registry.get_transport')
     @mock.patch('raven.base.ClientState.should_try')
     def test_async_send_remote_failover(self, should_try, get_transport):