Add heartbeat_timeout to the Client options.

This setting configures how long before a timeout event is emitted
internally and disconnects the websocket. Since some users were
experiencing issues with the gateway not responding, this should help
mitigate the issue for those with poor PCs.

4 files changed, 11 insertions, 2 deletions

diff --git a/discord/client.py b/discord/client.py
index 25fd5bb7..a310c3de 100644
--- a/discord/client.py
+++ b/discord/client.py
@@ -95,6 +95,11 @@ class Client:
         A game to start your presence with upon logging on to Discord.
     status: Optional[:class:`Status`]
         A status to start your presence with upon logging on to Discord.
+    heartbeat_timeout: float
+        The maximum numbers of seconds before timing out and restarting the
+        WebSocket in the case of not receiving a HEARTBEAT_ACK. Useful if
+        processing the initial packets take too long to the point of disconnecting
+        you. The default timeout is 60 seconds.
 
     Attributes
     -----------
diff --git a/discord/gateway.py b/discord/gateway.py
index d5ddb9eb..50104034 100644
--- a/discord/gateway.py
+++ b/discord/gateway.py
@@ -64,10 +64,11 @@ class KeepAliveHandler(threading.Thread):
         self._stop_ev = threading.Event()
         self._last_ack = time.time()
         self._last_send = time.time()
+        self.heartbeat_timeout = ws._max_heartbeat_timeout
 
     def run(self):
         while not self._stop_ev.wait(self.interval):
-            if self._last_ack + 2 * self.interval < time.time():
+            if self._last_ack + self.heartbeat_timeout < time.time():
                 log.warn("Shard ID %s has stopped responding to the gateway. Closing and restarting." % self.shard_id)
                 coro = self.ws.close(1006)
                 f = compat.run_coroutine_threadsafe(coro, loop=self.ws.loop)
@@ -205,6 +206,7 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
         ws.shard_count = client._connection.shard_count
         ws.session_id = session
         ws.sequence = sequence
+        ws._max_heartbeat_timeout = client._connection.heartbeat_timeout
 
         client._connection._update_references(ws)
 
@@ -590,6 +592,7 @@ class DiscordVoiceWebSocket(websockets.client.WebSocketClientProtocol):
         ws = yield from websockets.connect(gateway, loop=client.loop, klass=cls)
         ws.gateway = gateway
         ws._connection = client
+        ws._max_heartbeat_timeout = 60.0
 
         if resume:
             yield from ws.resume()
diff --git a/discord/shard.py b/discord/shard.py
index 75fd661e..70b99ee6 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -204,6 +204,7 @@ class AutoShardedClient(Client):
         ws.gateway = gateway
         ws.shard_id = shard_id
         ws.shard_count = self.shard_count
+        ws._max_heartbeat_timeout = self._connection.heartbeat_timeout
 
         try:
             # OP HELLO
diff --git a/discord/state.py b/discord/state.py
index a8d4ffe7..0511b754 100644
--- a/discord/state.py
+++ b/discord/state.py
@@ -64,6 +64,7 @@ class ConnectionState:
         self.shard_count = None
         self._ready_task = None
         self._fetch_offline = options.get('fetch_offline_members', True)
+        self.heartbeat_timeout = options.get('heartbeat_timeout', 60.0)
         self._listeners = []
 
         game = options.get('game', None)
@@ -907,7 +908,6 @@ class AutoShardedConnectionState(ConnectionState):
             launch.set()
             yield from asyncio.sleep(2.0 * self.shard_count, loop=self.loop)
 
-
         if self._fetch_offline:
             guilds = sorted(self._ready_state.guilds, key=lambda g: g.shard_id)