Handle rate limiting on RTM connections (#306)

* Added automatic retry for rtm.connect/rtm.start calls
* Improved error handling and added debug output

Author: Roach

slackclient/server.py

@@ -38,9 +38,10 @@ def __init__(self, token, connect=True, proxies=None):
         self.websocket = None
         self.ws_url = None
         self.connected = False
-        self.last_connected_at = 0
         self.auto_reconnect = False
-        self.reconnect_attempt = 0
+        self.last_connected_at = 0
+        self.reconnect_count = 0
+        self.rtm_connect_retries = 0
 
         # Connect to RTM on load
         if connect:
@@ -90,7 +91,7 @@ def rtm_connect(self, reconnect=False, timeout=None, use_rtm_start=True, **kwarg
 
         :Args:
             reconnect (boolean) Whether this method is being called to reconnect to RTM
-            timeout (int): Timeout for Web API calls
+            timeout (int): Stop waiting for Web API response after this many seconds
             use_rtm_start (boolean): `True` to connect using `rtm.start` or
             `False` to connect using`rtm.connect`
             https://api.slack.com/rtm#connecting_with_rtm.connect_vs._rtm.start
@@ -104,34 +105,42 @@ def rtm_connect(self, reconnect=False, timeout=None, use_rtm_start=True, **kwarg
         connect_method = "rtm.start" if use_rtm_start else "rtm.connect"
 
         # If the `auto_reconnect` param was passed, set the server's `auto_reconnect` attr
-        if kwargs and kwargs["auto_reconnect"] is True:
-            self.auto_reconnect = True
+        if 'auto_reconnect' in kwargs:
+            self.auto_reconnect = kwargs["auto_reconnect"]
 
         # If this is an auto reconnect, rate limit reconnect attempts
         if self.auto_reconnect and reconnect:
             # Raise a SlackConnectionError after 5 retries within 3 minutes
-            recon_attempt = self.reconnect_attempt
-            if recon_attempt == 5:
+            recon_count = self.reconnect_count
+            if recon_count == 5:
                 logging.error("RTM connection failed, reached max reconnects.")
                 raise SlackConnectionError("RTM connection failed, reached max reconnects.")
             # Wait to reconnect if the last reconnect was more than 3 minutes ago
             if (time.time() - self.last_connected_at) < 180:
-                if recon_attempt > 0:
+                if recon_count > 0:
                     # Back off after the the first attempt
                     backoff_offset_multiplier = random.randint(1, 4)
-                    retry_timeout = (backoff_offset_multiplier * recon_attempt * recon_attempt)
+                    retry_timeout = (backoff_offset_multiplier * recon_count * recon_count)
                     logging.debug("Reconnecting in %d seconds", retry_timeout)
 
                     time.sleep(retry_timeout)
-                self.reconnect_attempt += 1
+                self.reconnect_count += 1
             else:
-                self.reconnect_attempt = 0
+                self.reconnect_count = 0
 
         reply = self.api_requester.do(self.token, connect_method, timeout=timeout, post_data=kwargs)
 
         if reply.status_code != 200:
-            raise SlackConnectionError(reply=reply)
+            if self.rtm_connect_retries < 5 and reply.status_code == 429:
+                self.rtm_connect_retries += 1
+                retry_after = int(reply.headers.get('retry-after', 120))
+                logging.debug("HTTP 429: Rate limited. Retrying in %d seconds", retry_after)
+                time.sleep(retry_after)
+                self.rtm_connect(reconnect=reconnect, timeout=timeout)
+            else:
+                raise SlackConnectionError("RTM connection attempt was rate limited 10 times.")
         else:
+            self.rtm_connect_retries = 0
             login_data = reply.json()
             if login_data["ok"]:
                 self.ws_url = login_data['url']

test_requirements.txt

@@ -1,6 +1,7 @@
 pytest<3.3
 codecov
 coverage
+mock
 pytest-cov
 pytest-mock
 responses

tests/test_server.py

@@ -3,6 +3,7 @@
 import pytest
 import requests
 import responses
+from mock import patch
 from slackclient.user import User
 from slackclient.server import Server, SlackLoginError, SlackConnectionError
 from slackclient.channel import Channel
@@ -51,25 +52,23 @@ def test_server_is_hashable(server):
     assert (server_map[server] == 'foo') is False
 
 
-def test_response_headers(server):
+@patch('time.sleep', return_value=None)
+def test_rate_limiting(patched_time_sleep, server):
     # Testing for rate limit retry headers
     with responses.RequestsMock() as rsps:
         rsps.add(
             responses.POST,
-            "https://slack.com/api/auth.test",
+            "https://slack.com/api/rtm.start",
             status=429,
             json={"ok": False},
             headers={'Retry-After': "1"}
         )
 
-        res = json.loads(server.api_call("auth.test"))
-
-        for call in rsps.calls:
-            assert call.request.url in [
-                "https://slack.com/api/auth.test"
-            ]
-            assert call.response.status_code == 429
-        assert res["headers"]['Retry-After'] == "1"
+        with pytest.raises(SlackConnectionError) as e:
+            server.rtm_connect()
+            for call in rsps.calls:
+                assert call.response.status_code == 429
+            assert e.message == "RTM connection attempt was rate limited 10 times."
 
 
 def test_custom_agent(server):
@@ -147,6 +146,23 @@ def test_rtm_reconnect(server, rtm_start_fixture):
             ]
 
 
+@patch('time.sleep', return_value=None)
+def test_rtm_max_reconnect_timeout(patched_time_sleep, server, rtm_start_fixture):
+    with responses.RequestsMock() as rsps:
+        rsps.add(
+            responses.POST,
+            "https://slack.com/api/rtm.connect",
+            status=200,
+            json=rtm_start_fixture
+        )
+
+        server.reconnect_count = 4
+        server.last_connected_at = time.time()
+        server.rtm_connect(auto_reconnect=True, reconnect=True, use_rtm_start=False)
+
+        assert server.reconnect_count == 5
+
+
 def test_rtm_reconnect_timeout_recently_connected(server, rtm_start_fixture):
     # If reconnected recently, server must wait to reconnect and increment the counter
     with responses.RequestsMock() as rsps:
@@ -157,11 +173,11 @@ def test_rtm_reconnect_timeout_recently_connected(server, rtm_start_fixture):
             json=rtm_start_fixture
         )
 
-        server.reconnect_attempt = 0
+        server.reconnect_count = 0
         server.last_connected_at = time.time()
         server.rtm_connect(auto_reconnect=True, reconnect=True, use_rtm_start=False)
 
-        assert server.reconnect_attempt == 1
+        assert server.reconnect_count == 1
         for call in rsps.calls:
             assert call.request.url in [
                 "https://slack.com/api/rtm.connect"
@@ -178,20 +194,21 @@ def test_rtm_reconnect_timeout_not_recently_connected(server, rtm_start_fixture)
             json=rtm_start_fixture
         )
 
-        server.reconnect_attempt = 1
+        server.reconnect_count = 1
         server.last_connected_at = time.time() - 180
         server.rtm_connect(auto_reconnect=True, reconnect=True, use_rtm_start=False)
 
-        assert server.reconnect_attempt == 0
+        assert server.reconnect_count == 0
         for call in rsps.calls:
             assert call.request.url in [
                 "https://slack.com/api/rtm.connect"
             ]
 
 
-def test_max_rtm_reconnect_attempts(server):
+def test_max_rtm_reconnects(server, monkeypatch):
+    monkeypatch.setattr("time.sleep", None)
     with pytest.raises(SlackConnectionError) as e:
-        server.reconnect_attempt = 5
+        server.reconnect_count = 5
         server.rtm_connect(auto_reconnect=True, reconnect=True, use_rtm_start=False)
         assert e.message == "RTM connection failed, reached max reconnects."