feat: add shard information to all state & gw logging (#1551)

* feat: add shard information to all state & gw logging

* feat: restore logging's heart

Author: LordOfPolls

interactions/api/gateway/gateway.py

@@ -1,5 +1,6 @@
 """Outlines the interaction between interactions and Discord's Gateway API."""
 import asyncio
+import logging
 import sys
 import time
 import zlib
@@ -171,31 +172,32 @@ async def run(self) -> None:
     async def dispatch_opcode(self, data, op: OPCODE) -> None:
         match op:
             case OPCODE.HEARTBEAT:
-                self.logger.debug("Received heartbeat request from gateway")
+                self.state.wrapped_logger(logging.DEBUG, "❤ Received heartbeat request from gateway")
                 return await self.send_heartbeat()
 
             case OPCODE.HEARTBEAT_ACK:
                 self._latency.append(time.perf_counter() - self._last_heartbeat)
 
                 if self._last_heartbeat != 0 and self._latency[-1] >= 15:
-                    self.logger.warning(
-                        f"High Latency! shard ID {self.shard[0]} heartbeat took {self._latency[-1]:.1f}s to be acknowledged!"
+                    self.state.wrapped_logger(
+                        logging.WARNING,
+                        f"❤ High Latency! shard ID {self.shard[0]} heartbeat took {self._latency[-1]:.1f}s to be acknowledged!",
                     )
                 else:
-                    self.logger.debug(f"❤ Heartbeat acknowledged after {self._latency[-1]:.5f} seconds")
+                    self.state.wrapped_logger(logging.DEBUG, "❤ Received heartbeat acknowledgement from gateway")
 
                 return self._acknowledged.set()
 
             case OPCODE.RECONNECT:
-                self.logger.debug("Gateway requested reconnect. Reconnecting...")
+                self.state.wrapped_logger(logging.DEBUG, "Gateway requested reconnect. Reconnecting...")
                 return await self.reconnect(resume=True, url=self.ws_resume_url)
 
             case OPCODE.INVALIDATE_SESSION:
-                self.logger.warning("Gateway has invalidated session! Reconnecting...")
+                self.state.wrapped_logger(logging.WARNING, "Gateway invalidated session. Reconnecting...")
                 return await self.reconnect()
 
             case _:
-                return self.logger.debug(f"Unhandled OPCODE: {op} = {OPCODE(op).name}")
+                return self.state.wrapped_logger(logging.DEBUG, f"Unhandled OPCODE: {op} = {OPCODE(op).name}")
 
     async def dispatch_event(self, data, seq, event) -> None:
         match event:
@@ -207,12 +209,14 @@ async def dispatch_event(self, data, seq, event) -> None:
                 self.ws_resume_url = (
                     f"{data['resume_gateway_url']}?encoding=json&v={__api_version__}&compress=zlib-stream"
                 )
-                self.logger.info(f"Shard {self.shard[0]} has connected to gateway!")
-                self.logger.debug(f"Session ID: {self.session_id} Trace: {self._trace}")
+                self.state.wrapped_logger(logging.INFO, "Gateway connection established")
+                self.state.wrapped_logger(logging.DEBUG, f"Session ID: {self.session_id} Trace: {self._trace}")
                 return self.state.client.dispatch(events.WebsocketReady(data))
 
             case "RESUMED":
-                self.logger.info(f"Successfully resumed connection! Session_ID: {self.session_id}")
+                self.state.wrapped_logger(
+                    logging.INFO, f"Successfully resumed connection! Session_ID: {self.session_id}"
+                )
                 self.state.client.dispatch(events.Resume())
                 return None
 
@@ -228,9 +232,11 @@ async def dispatch_event(self, data, seq, event) -> None:
                             processor(events.RawGatewayEvent(data.copy(), override_name=event_name))
                         )
                     except Exception as ex:
-                        self.logger.error(f"Failed to run event processor for {event_name}: {ex}")
+                        self.state.wrapped_logger(
+                            logging.ERROR, f"Failed to run event processor for {event_name}: {ex}"
+                        )
                 else:
-                    self.logger.debug(f"No processor for `{event_name}`")
+                    self.state.wrapped_logger(logging.DEBUG, f"No processor for `{event_name}`")
 
         self.state.client.dispatch(events.RawGatewayEvent(data.copy(), override_name="raw_gateway_event"))
         self.state.client.dispatch(events.RawGatewayEvent(data.copy(), override_name=f"raw_{event.lower()}"))
@@ -263,8 +269,8 @@ async def _identify(self) -> None:
         serialized = FastJson.dumps(payload)
         await self.ws.send_str(serialized)
 
-        self.logger.debug(
-            f"Shard ID {self.shard[0]} has identified itself to Gateway, requesting intents: {self.state.intents}!"
+        self.state.wrapped_logger(
+            logging.DEBUG, f"Identification payload sent to gateway, requesting intents: {self.state.intents}"
         )
 
     async def reconnect(self, *, resume: bool = False, code: int = 1012, url: str | None = None) -> None:
@@ -289,11 +295,11 @@ async def _resume_connection(self) -> None:
         serialized = FastJson.dumps(payload)
         await self.ws.send_str(serialized)
 
-        self.logger.debug(f"{self.shard[0]} is attempting to resume a connection")
+        self.state.wrapped_logger(logging.DEBUG, f"Resume payload sent to gateway, session ID: {self.session_id}")
 
     async def send_heartbeat(self) -> None:
         await self.send_json({"op": OPCODE.HEARTBEAT, "d": self.sequence}, bypass=True)
-        self.logger.debug(f"❤ Shard {self.shard[0]} is sending a Heartbeat")
+        self.state.wrapped_logger(logging.DEBUG, "❤ Gateway is sending a Heartbeat")
 
     async def change_presence(self, activity=None, status: Status = Status.ONLINE, since=None) -> None:
         """Update the bot's presence status."""

interactions/api/gateway/state.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 import traceback
 from datetime import datetime
 from logging import Logger
@@ -72,7 +73,7 @@ async def start(self) -> None:
         """Connect to the Discord Gateway."""
         self.gateway_url = await self.client.http.get_gateway()
 
-        self.logger.debug(f"Starting Shard ID {self.shard_id}")
+        self.wrapped_logger(logging.INFO, "Starting Shard")
         self.start_time = datetime.now()
         self._shard_task = asyncio.create_task(self._ws_connect())
 
@@ -84,7 +85,7 @@ async def start(self) -> None:
 
     async def stop(self) -> None:
         """Disconnect from the Discord Gateway."""
-        self.logger.debug(f"Shutting down shard ID {self.shard_id}")
+        self.wrapped_logger(logging.INFO, "Stopping Shard")
         if self.gateway is not None:
             self.gateway.close()
             self.gateway = None
@@ -102,7 +103,7 @@ def clear_ready(self) -> None:
 
     async def _ws_connect(self) -> None:
         """Connect to the Discord Gateway."""
-        self.logger.info(f"Shard {self.shard_id} is attempting to connect to gateway...")
+        self.wrapped_logger(logging.INFO, "Shard is attempting to connect to gateway...")
         try:
             async with GatewayClient(self, (self.shard_id, self.client.total_shards)) as self.gateway:
                 try:
@@ -127,7 +128,18 @@ async def _ws_connect(self) -> None:
 
         except Exception as e:
             self.client.dispatch(events.Disconnect())
-            self.logger.error("".join(traceback.format_exception(type(e), e, e.__traceback__)))
+            self.wrapped_logger("".join(traceback.format_exception(type(e), e, e.__traceback__)))
+
+    def wrapped_logger(self, level: int, message: str, **kwargs) -> None:
+        """
+        A logging wrapper that adds shard information to the message.
+
+        Args:
+            level: The logging level
+            message: The message to log
+            **kwargs: Any additional keyword arguments that Logger.log accepts
+        """
+        self.logger.log(level, f"Shard ID {self.shard_id} | {message}", **kwargs)
 
     async def change_presence(
         self,
@@ -157,15 +169,19 @@ async def change_presence(
 
             if activity.type == ActivityType.STREAMING:
                 if not activity.url:
-                    self.logger.warning("Streaming activity cannot be set without a valid URL attribute")
+                    self.wrapped_logger(
+                        logging.WARNING, "Streaming activity cannot be set without a valid URL attribute"
+                    )
             elif activity.type not in [
                 ActivityType.GAME,
                 ActivityType.STREAMING,
                 ActivityType.LISTENING,
                 ActivityType.WATCHING,
                 ActivityType.COMPETING,
             ]:
-                self.logger.warning(f"Activity type `{ActivityType(activity.type).name}` may not be enabled for bots")
+                self.wrapped_logger(
+                    logging.WARNING, f"Activity type `{ActivityType(activity.type).name}` may not be enabled for bots"
+                )
         if status:
             if not isinstance(status, Status):
                 try:
@@ -175,7 +191,7 @@ async def change_presence(
         elif self.client.status:
             status = self.client.status
         else:
-            self.logger.warning("Status must be set to a valid status type, defaulting to online")
+            self.wrapped_logger(logging.WARNING, "Status must be set to a valid status type, defaulting to online")
             status = Status.ONLINE
 
         self.client._status = status