Add enable_debug_events parameter to enable expensive debug events

This should allow less dispatching during heavy gateway event streams.

2 files changed, 26 insertions, 3 deletions

diff --git a/discord/client.py b/discord/client.py
index 357c1ab6..a41e061c 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 163f480a..e0814935 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))