Browse Source

Add RESUME support.

pull/237/head
Rapptz 9 years ago
parent
commit
e0a91df32b
  1. 7
      discord/client.py
  2. 52
      discord/gateway.py
  3. 3
      discord/state.py
  4. 4
      docs/api.rst

7
discord/client.py

@ -401,9 +401,10 @@ class Client:
while not self.is_closed: while not self.is_closed:
try: try:
yield from self.ws.poll_event() yield from self.ws.poll_event()
except ReconnectWebSocket: except (ReconnectWebSocket, ResumeWebSocket) as e:
log.info('Reconnecting the websocket.') resume = type(e) is ResumeWebSocket
self.ws = yield from DiscordWebSocket.from_client(self) log.info('Got ' + type(e).__name__)
self.ws = yield from DiscordWebSocket.from_client(self, resume=resume)
except ConnectionClosed as e: except ConnectionClosed as e:
yield from self.close() yield from self.close()
if e.code != 1000: if e.code != 1000:

52
discord/gateway.py

@ -42,12 +42,16 @@ log = logging.getLogger(__name__)
__all__ = [ 'ReconnectWebSocket', 'get_gateway', 'DiscordWebSocket', __all__ = [ 'ReconnectWebSocket', 'get_gateway', 'DiscordWebSocket',
'KeepAliveHandler', 'VoiceKeepAliveHandler', 'KeepAliveHandler', 'VoiceKeepAliveHandler',
'DiscordVoiceWebSocket' ] 'DiscordVoiceWebSocket', 'ResumeWebSocket' ]
class ReconnectWebSocket(Exception): class ReconnectWebSocket(Exception):
"""Signals to handle the RECONNECT opcode.""" """Signals to handle the RECONNECT opcode."""
pass pass
class ResumeWebSocket(Exception):
"""Signals to initialise via RESUME opcode instead of IDENTIFY."""
pass
EventListener = namedtuple('EventListener', 'predicate event result future') EventListener = namedtuple('EventListener', 'predicate event result future')
class KeepAliveHandler(threading.Thread): class KeepAliveHandler(threading.Thread):
@ -179,10 +183,9 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
# the keep alive # the keep alive
self._keep_alive = None self._keep_alive = None
@classmethod @classmethod
@asyncio.coroutine @asyncio.coroutine
def from_client(cls, client): def from_client(cls, client, *, resume=False):
"""Creates a main websocket for Discord from a :class:`Client`. """Creates a main websocket for Discord from a :class:`Client`.
This is for internal use only. This is for internal use only.
@ -197,9 +200,21 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
ws.gateway = gateway ws.gateway = gateway
log.info('Created websocket connected to {}'.format(gateway)) log.info('Created websocket connected to {}'.format(gateway))
yield from ws.identify() if not resume:
log.info('sent the identify payload to create the websocket') yield from ws.identify()
return ws 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): def wait_for(self, event, predicate, result=None):
"""Waits for a DISPATCH'd event that meets the predicate. """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) 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 @asyncio.coroutine
def received_message(self, msg): def received_message(self, msg):
self._dispatch('socket_raw_receive', msg) self._dispatch('socket_raw_receive', msg)
@ -271,13 +301,14 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
# "reconnect" can only be handled by the Client # "reconnect" can only be handled by the Client
# so we terminate our connection and raise an # so we terminate our connection and raise an
# internal exception signalling to reconnect. # internal exception signalling to reconnect.
log.info('Receivede RECONNECT opcode.') log.info('Received RECONNECT opcode.')
yield from self.close() yield from self.close()
raise ReconnectWebSocket() raise ReconnectWebSocket()
if op == self.INVALIDATE_SESSION: if op == self.INVALIDATE_SESSION:
state.sequence = None state.sequence = None
state.session_id = None state.session_id = None
yield from self.identify()
return return
if op != self.DISPATCH: if op != self.DISPATCH:
@ -347,8 +378,11 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol):
yield from self.received_message(msg) yield from self.received_message(msg)
except websockets.exceptions.ConnectionClosed as e: except websockets.exceptions.ConnectionClosed as e:
if self._can_handle_close(e.code): if self._can_handle_close(e.code):
log.info('Websocket closed with {0.code}, attempting a reconnect.'.format(e)) log.info('Websocket closed with {0.code} ({0.reason}), attempting a reconnect.'.format(e))
raise ReconnectWebSocket() from e if e.code == 4006:
raise ReconnectWebSocket() from e
else:
raise ResumeWebSocket() from e
else: else:
raise ConnectionClosed(e) from e raise ConnectionClosed(e) from e

3
discord/state.py

@ -199,6 +199,9 @@ class ConnectionState:
compat.create_task(self._delay_ready(), loop=self.loop) compat.create_task(self._delay_ready(), loop=self.loop)
def parse_resumed(self, data):
self.dispatch('resumed')
def parse_message_create(self, data): def parse_message_create(self, data):
channel = self.get_channel(data.get('channel_id')) channel = self.get_channel(data.get('channel_id'))
message = Message(channel=channel, **data) message = Message(channel=channel, **data)

4
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. 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) .. function:: on_error(event, \*args, \*\*kwargs)
Usually when an event raises an uncaught exception, a traceback is Usually when an event raises an uncaught exception, a traceback is

Loading…
Cancel
Save