Bump waiting time of GUILD_CREATE stream by number of shards waited.

2 seconds might be too short and cause it to finish waiting while other
shards are still IDENTIFYing.

2 files changed, 4 insertions, 2 deletions

diff --git a/discord/shard.py b/discord/shard.py
index 237e50e5..53c09c0d 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -251,6 +251,7 @@ class AutoShardedClient(Client):
         self._connection.shard_count = self.shard_count
 
         shard_ids = self.shard_ids if self.shard_ids else range(self.shard_count)
+        self._connection.shard_ids = shard_ids
 
         for shard_id in shard_ids:
             await self.launch_shard(gateway, shard_id)
diff --git a/discord/state.py b/discord/state.py
index 738495a2..e4ecdc45 100644
--- a/discord/state.py
+++ b/discord/state.py
@@ -1029,6 +1029,7 @@ class AutoShardedConnectionState(ConnectionState):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._ready_task = None
+        self.shard_ids = ()
 
     async def request_offline_members(self, guilds, *, shard_id):
         # get all the chunks
@@ -1051,10 +1052,10 @@ class AutoShardedConnectionState(ConnectionState):
     async def _delay_ready(self):
         launch = self._ready_state.launch
         while True:
-            # this snippet of code is basically waiting 2 seconds
+            # this snippet of code is basically waiting 2 * shard_ids seconds
             # until the last GUILD_CREATE was sent
             try:
-                await asyncio.wait_for(launch.wait(), timeout=2.0)
+                await asyncio.wait_for(launch.wait(), timeout=2.0 * len(self.shard_ids))
             except asyncio.TimeoutError:
                 break
             else: