From b6404933001207dcd03b59f185cce0221b071147 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Sat, 31 Jul 2021 20:35:16 -0400 Subject: [PATCH] Add enable_debug_events parameter to enable expensive debug events This should allow less dispatching during heavy gateway event streams. --- discord/client.py | 9 +++++++++ discord/gateway.py | 20 +++++++++++++++++--- docs/api.rst | 4 ++++ 3 files changed, 30 insertions(+), 3 deletions(-) diff --git a/discord/client.py b/discord/client.py index 357c1ab66..a41e061c6 100644 --- a/discord/client.py +++ b/discord/client.py @@ -185,6 +185,14 @@ class Client: sync your system clock to Google's NTP server. .. versionadded:: 1.3 + enable_debug_events: :class:`bool` + Whether to enable events that are useful only for debugging gateway related information. + + Right now this involves :func:`on_socket_raw_receive` and :func`:`on_socket_raw_send`. If + this is ``False`` then those events will not be dispatched (due to performance considerations). + To enable these events, this must be set to ``True``. Defaults to ``False``. + + .. versionadded:: 2.0 Attributes ----------- @@ -219,6 +227,7 @@ class Client: 'before_identify': self._call_before_identify_hook } + self._enable_debug_events: bool = options.pop('enable_debug_events', False) self._connection: ConnectionState = self._get_state(**options) self._connection.shard_count = self.shard_count self._closed: bool = False diff --git a/discord/gateway.py b/discord/gateway.py index 163f480a2..e08149357 100644 --- a/discord/gateway.py +++ b/discord/gateway.py @@ -293,6 +293,12 @@ class DiscordWebSocket: def is_ratelimited(self): return self._rate_limiter.is_ratelimited() + def log_receive(self, data, /): + self._dispatch('socket_raw_receive', data) + + def empty_log_receive(self, _, /): + pass + @classmethod async def from_client(cls, client, *, initial=False, gateway=None, shard_id=None, session=None, sequence=None, resume=False): """Creates a main websocket for Discord from a :class:`Client`. @@ -318,6 +324,10 @@ class DiscordWebSocket: ws.sequence = sequence ws._max_heartbeat_timeout = client._connection.heartbeat_timeout + if client._enable_debug_events: + ws.send = ws.debug_send + ws.log_receive = ws.empty_log_receive + client._connection._update_references(ws) log.debug('Created websocket connected to %s', gateway) @@ -409,8 +419,8 @@ class DiscordWebSocket: await self.send_as_json(payload) log.info('Shard ID %s has sent the RESUME payload.', self.shard_id) - async def received_message(self, msg): - self._dispatch('socket_raw_receive', msg) + async def received_message(self, msg, /): + self.log_receive(msg) if type(msg) is bytes: self._buffer.extend(msg) @@ -574,11 +584,15 @@ class DiscordWebSocket: log.info('Websocket closed with %s, cannot reconnect.', code) raise ConnectionClosed(self.socket, shard_id=self.shard_id, code=code) from None - async def send(self, data): + async def debug_send(self, data, /): await self._rate_limiter.block() self._dispatch('socket_raw_send', data) await self.socket.send_str(data) + async def send(self, data, /): + await self._rate_limiter.block() + await self.socket.send_str(data) + async def send_as_json(self, data): try: await self.send(utils.to_json(data)) diff --git a/docs/api.rst b/docs/api.rst index 77be49ca6..c8a09d881 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -318,6 +318,8 @@ to handle it, which defaults to print a traceback and ignoring the exception. This is only really useful for grabbing the WebSocket stream and debugging purposes. + This requires setting the ``enable_debug_events`` setting in the :class:`Client`. + .. note:: This is only for the messages received from the client @@ -337,6 +339,8 @@ to handle it, which defaults to print a traceback and ignoring the exception. This is only really useful for grabbing the WebSocket stream and debugging purposes. + This requires setting the ``enable_debug_events`` setting in the :class:`Client`. + .. note:: This is only for the messages sent from the client