mirrors/discord.py
1
0
Fork 0
mirror of https://github.com/Rapptz/discord.py.git synced 2025-06-07 20:28:38 +00:00
Code Issues Projects Releases Wiki Activity

Add a timeout for receiving websocket messages.

Browse Source
This commit is contained in:
Rapptz 2020-05-24 04:05:09 -04:00
parent 137ff34f3a
commit e0660ef8a6
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
discord/gateway.py
View File

@ -96,7 +96,7 @@ class KeepAliveHandler(threading.Thread):
try: try:
f.result() f.result()
except Exception: except Exception:
pass log.exception('An error occurred while stopping the gateway. Ignoring.')
finally: finally:
self.stop() self.stop()
return return
@ -500,7 +500,7 @@ class DiscordWebSocket:
The websocket connection was terminated for unhandled reasons. The websocket connection was terminated for unhandled reasons.
""" """
try: try:
msg = await self.socket.receive() msg = await self.socket.receive(timeout=self._max_heartbeat_timeout)
if msg.type is aiohttp.WSMsgType.TEXT: if msg.type is aiohttp.WSMsgType.TEXT:
await self.received_message(msg.data) await self.received_message(msg.data)
elif msg.type is aiohttp.WSMsgType.BINARY: elif msg.type is aiohttp.WSMsgType.BINARY:
@ -511,12 +511,16 @@ class DiscordWebSocket:
elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.CLOSE): elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.CLOSE):
log.debug('Received %s', msg) log.debug('Received %s', msg)
raise WebSocketClosure raise WebSocketClosure
except WebSocketClosure: except (asyncio.TimeoutError, WebSocketClosure) as e:
# Ensure the keep alive handler is closed # Ensure the keep alive handler is closed
if self._keep_alive: if self._keep_alive:
self._keep_alive.stop() self._keep_alive.stop()
self._keep_alive = None self._keep_alive = None
if isinstance(e, asyncio.TimeoutError):
log.info('Timed out receiving packet. Attempting a reconnect.')
raise ReconnectWebSocket(self.shard_id) from None
if self._can_handle_close(): if self._can_handle_close():
log.info('Websocket closed with %s, attempting a reconnect.', self.socket.close_code) log.info('Websocket closed with %s, attempting a reconnect.', self.socket.close_code)
raise ReconnectWebSocket(self.shard_id) from None raise ReconnectWebSocket(self.shard_id) from None
@ -819,7 +823,7 @@ class DiscordVoiceWebSocket:
elif msg.type is aiohttp.WSMsgType.ERROR: elif msg.type is aiohttp.WSMsgType.ERROR:
log.debug('Received %s', msg) log.debug('Received %s', msg)
raise ConnectionClosed(self.ws, shard_id=None) from msg.data raise ConnectionClosed(self.ws, shard_id=None) from msg.data
elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE): elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING):
log.debug('Received %s', msg) log.debug('Received %s', msg)
raise ConnectionClosed(self.ws, shard_id=None) raise ConnectionClosed(self.ws, shard_id=None)