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)