From de65f7309bc8f8c931f29dcd76ebd9e8b0ee547b Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Tue, 8 Aug 2017 21:12:04 -0400
Subject: [PATCH] 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.
---
 discord/client.py  | 5 +++++
 discord/gateway.py | 5 ++++-
 discord/shard.py   | 1 +
 discord/state.py   | 2 +-
 4 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/discord/client.py b/discord/client.py
index 25fd5bb78..a310c3deb 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 d5ddb9eb8..50104034f 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 75fd661e0..70b99ee6b 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 a8d4ffe75..0511b7545 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)