From e0a91df32beabafa990ce09dee13af665c770079 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Wed, 1 Jun 2016 05:13:15 -0400 Subject: [PATCH] Add RESUME support. --- discord/client.py | 7 ++++--- discord/gateway.py | 52 ++++++++++++++++++++++++++++++++++++++-------- discord/state.py | 3 +++ docs/api.rst | 4 ++++ 4 files changed, 54 insertions(+), 12 deletions(-) diff --git a/discord/client.py b/discord/client.py index 7ac7ca66d..295768455 100644 --- a/discord/client.py +++ b/discord/client.py @@ -401,9 +401,10 @@ class Client: while not self.is_closed: try: yield from self.ws.poll_event() - except ReconnectWebSocket: - log.info('Reconnecting the websocket.') - self.ws = yield from DiscordWebSocket.from_client(self) + except (ReconnectWebSocket, ResumeWebSocket) as e: + resume = type(e) is ResumeWebSocket + log.info('Got ' + type(e).__name__) + self.ws = yield from DiscordWebSocket.from_client(self, resume=resume) except ConnectionClosed as e: yield from self.close() if e.code != 1000: diff --git a/discord/gateway.py b/discord/gateway.py index 6261274d8..928e37be8 100644 --- a/discord/gateway.py +++ b/discord/gateway.py @@ -42,12 +42,16 @@ log = logging.getLogger(__name__) __all__ = [ 'ReconnectWebSocket', 'get_gateway', 'DiscordWebSocket', 'KeepAliveHandler', 'VoiceKeepAliveHandler', - 'DiscordVoiceWebSocket' ] + 'DiscordVoiceWebSocket', 'ResumeWebSocket' ] class ReconnectWebSocket(Exception): """Signals to handle the RECONNECT opcode.""" pass +class ResumeWebSocket(Exception): + """Signals to initialise via RESUME opcode instead of IDENTIFY.""" + pass + EventListener = namedtuple('EventListener', 'predicate event result future') class KeepAliveHandler(threading.Thread): @@ -179,10 +183,9 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): # the keep alive self._keep_alive = None - @classmethod @asyncio.coroutine - def from_client(cls, client): + def from_client(cls, client, *, resume=False): """Creates a main websocket for Discord from a :class:`Client`. This is for internal use only. @@ -197,9 +200,21 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): ws.gateway = gateway log.info('Created websocket connected to {}'.format(gateway)) - yield from ws.identify() - log.info('sent the identify payload to create the websocket') - return ws + if not resume: + yield from ws.identify() + log.info('sent the identify payload to create the websocket') + return ws + + yield from ws.resume() + log.info('sent the resume payload to create the websocket') + try: + yield from ws.ensure_open() + except websockets.exceptions.ConnectionClosed: + # ws got closed so let's just do a regular IDENTIFY connect. + log.info('RESUME failure.') + return (yield from cls.from_client(client)) + else: + return ws def wait_for(self, event, predicate, result=None): """Waits for a DISPATCH'd event that meets the predicate. @@ -247,6 +262,21 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): } yield from self.send_as_json(payload) + @asyncio.coroutine + def resume(self): + """Sends the RESUME packet.""" + state = self._connection + payload = { + 'op': self.RESUME, + 'd': { + 'seq': state.sequence, + 'session_id': state.session_id, + 'token': self.token + } + } + + yield from self.send_as_json(payload) + @asyncio.coroutine def received_message(self, msg): self._dispatch('socket_raw_receive', msg) @@ -271,13 +301,14 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): # "reconnect" can only be handled by the Client # so we terminate our connection and raise an # internal exception signalling to reconnect. - log.info('Receivede RECONNECT opcode.') + log.info('Received RECONNECT opcode.') yield from self.close() raise ReconnectWebSocket() if op == self.INVALIDATE_SESSION: state.sequence = None state.session_id = None + yield from self.identify() return if op != self.DISPATCH: @@ -347,8 +378,11 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): yield from self.received_message(msg) except websockets.exceptions.ConnectionClosed as e: if self._can_handle_close(e.code): - log.info('Websocket closed with {0.code}, attempting a reconnect.'.format(e)) - raise ReconnectWebSocket() from e + log.info('Websocket closed with {0.code} ({0.reason}), attempting a reconnect.'.format(e)) + if e.code == 4006: + raise ReconnectWebSocket() from e + else: + raise ResumeWebSocket() from e else: raise ConnectionClosed(e) from e diff --git a/discord/state.py b/discord/state.py index 7c1ba9937..634a2749b 100644 --- a/discord/state.py +++ b/discord/state.py @@ -199,6 +199,9 @@ class ConnectionState: compat.create_task(self._delay_ready(), loop=self.loop) + def parse_resumed(self, data): + self.dispatch('resumed') + def parse_message_create(self, data): channel = self.get_channel(data.get('channel_id')) message = Message(channel=channel, **data) diff --git a/docs/api.rst b/docs/api.rst index 91d66ab26..d1824ad2d 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -109,6 +109,10 @@ to handle it, which defaults to print a traceback and ignore the exception. This function is not guaranteed to be the first event called. +.. function:: on_resumed() + + Called when the client has resumed a session. + .. function:: on_error(event, \*args, \*\*kwargs) Usually when an event raises an uncaught exception, a traceback is