Rewrite on_ready delay to actually make sense and prevent heavy sleep

1 files changed, 15 insertions, 7 deletions

diff --git a/discord/state.py b/discord/state.py
index 5c9f7df4..0198414e 100644
--- a/discord/state.py
+++ b/discord/state.py
@@ -340,11 +340,15 @@ class ConnectionState:
 
             # only real bots wait for GUILD_CREATE streaming
             if self.is_bot:
-                while not launch.is_set():
+                while True:
                     # this snippet of code is basically waiting 2 seconds
                     # until the last GUILD_CREATE was sent
-                    launch.set()
-                    await asyncio.sleep(2)
+                    try:
+                        await asyncio.wait_for(launch.wait(), timeout=2.0)
+                    except asyncio.TimeoutError:
+                        break
+                    else:
+                        launch.clear()
 
             guilds = next(zip(*self._ready_state.guilds), [])
             if self._fetch_offline:
@@ -729,7 +733,7 @@ class ConnectionState:
                 # so we say.
                 try:
                     state = self._ready_state
-                    state.launch.clear()
+                    state.launch.set()
                     state.guilds.append((guild, unavailable))
                 except AttributeError:
                     # the _ready_state attribute is only there during
@@ -1020,11 +1024,15 @@ class AutoShardedConnectionState(ConnectionState):
 
     async def _delay_ready(self):
         launch = self._ready_state.launch
-        while not launch.is_set():
+        while True:
             # this snippet of code is basically waiting 2 seconds
             # until the last GUILD_CREATE was sent
-            launch.set()
-            await asyncio.sleep(2.0 * self.shard_count)
+            try:
+                await asyncio.wait_for(launch.wait(), timeout=2.0)
+            except asyncio.TimeoutError:
+                break
+            else:
+                launch.clear()
 
         guilds = sorted(self._ready_state.guilds, key=lambda g: g[0].shard_id)