Handle waiting for a ready within our own connection

Voxel-Fox-Ltd · Jan 17, 2024 · 956bd83 · 956bd83
1 parent 282db8f
commit 956bd83
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/novus/api/gateway/dispatch.py b/novus/api/gateway/dispatch.py
@@ -160,13 +160,9 @@ async def handle_dispatch(self, event_name: str, data: dict) -> None:
             self.cache.application_id = try_snowflake(data['application']['id'])
             self.shard.resume_url = data['resume_gateway_url']
             self.shard.session_id = data['session_id']
-            self.shard.ready.set()
-            log.info("[%s] Received ready", self.shard.shard_id)
             self.dispatch("READY")
             return None
         elif event_name == "RESUMED":
-            self.shard.ready.set()
-            log.info("Received resumed (shard %s)", self.shard.shard_id)
             return None
 
         coro: Callable[..., Awaitable[None]] | None

diff --git a/novus/api/gateway/gateway.py b/novus/api/gateway/gateway.py
@@ -524,17 +524,18 @@ async def _connect(
         self.socket = ws
 
         # Get hello
+        log.info("[%s] Waiting for a HELLO", self.shard_id)
         try:
             got = await asyncio.wait_for(self.receive(), timeout=10.0)
         except Exception as e:
             if attempt >= 5:
                 log.info(
-                    "[%s] Failed to connect to gateway (%s), closing (%s)",
+                    "[%s] Failed to get a HELLO (%s), closing (%s)",
                     self.shard_id, e, attempt,
                 )
                 return await self.close()
             log.info(
-                "[%s] Failed to connect to gateway (%s), reattempting (%s)",
+                "[%s] Failed to get a HELLO (%s), reattempting (%s)",
                 self.shard_id, e, attempt,
             )
             return await self._connect(
@@ -573,6 +574,7 @@ async def _connect(
                 reconnect=reconnect,
                 attempt=attempt + 1,
             )
+        log.info("[%s] Waiting for a READY", self.shard_id)
         try:
             await asyncio.wait_for(self.ready.wait(), timeout=30)
         except asyncio.TimeoutError:
@@ -790,6 +792,9 @@ async def message_handler(self) -> None:
                 case GatewayOpcode.dispatch:
                     event_name = cast(str, event_name)
                     sequence = cast(int, sequence)
+                    if event_name == "READY" or event_name == "RESUMED":
+                        log.info("[%s] Received %s", self.shard_id, event_name)
+                        self.ready.set()
                     self.sequence = sequence
                     t = asyncio.create_task(
                         self.handle_dispatch(event_name, message),