matthew/enhanced-discord.py
1
0
Fork 0
Code Issues 4 Pull Requests 21 Projects Releases Wiki Activity

Maximize concurrency when chunking on AutoSharded clients

Browse Source
This commit is contained in:
Rapptz 2020-09-05 13:48:02 -04:00
parent af8fc32329
commit 92e1816114
1 changed files with 45 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
discord/state.py
View File

@ -347,6 +347,7 @@ class ConnectionState:
try: try:
# only real bots wait for GUILD_CREATE streaming # only real bots wait for GUILD_CREATE streaming
if self.is_bot: if self.is_bot:
states = []
while True: while True:
# this snippet of code is basically waiting N seconds # this snippet of code is basically waiting N seconds
# until the last GUILD_CREATE was sent # until the last GUILD_CREATE was sent
@ -355,17 +356,26 @@ class ConnectionState:
except asyncio.TimeoutError: except asyncio.TimeoutError:
break break
else: else:
try: if self._fetch_offline:
if self._fetch_offline: future = await self.chunk_guild(guild, wait=False)
await asyncio.wait_for(self.chunk_guild(guild), timeout=60.0) states.append((guild, future))
except asyncio.TimeoutError: else:
log.info('Timed out waiting for chunks while launching ready event.')
finally:
if guild.unavailable is False: if guild.unavailable is False:
self.dispatch('guild_available', guild) self.dispatch('guild_available', guild)
else: else:
self.dispatch('guild_join', guild) self.dispatch('guild_join', guild)
for guild, future in states:
try:
await asyncio.wait_for(future, timeout=5.0)
except asyncio.TimeoutError:
log.warning('Shard ID %s timed out waiting for chunks for guild_id %s.', guild.shard_id, guild.id)
if guild.unavailable is False:
self.dispatch('guild_available', guild)
else:
self.dispatch('guild_join', guild)
# remove the state # remove the state
try: try:
del self._ready_state del self._ready_state
@ -733,12 +743,14 @@ class ConnectionState:
return self._add_guild_from_data(data) return self._add_guild_from_data(data)
async def chunk_guild(self, guild): async def chunk_guild(self, guild, *, wait=True):
future = self.loop.create_future() future = self.loop.create_future()
request = ChunkRequest(guild.id, future) request = ChunkRequest(guild.id, future)
self._chunk_requests.append(request) self._chunk_requests.append(request)
await self.chunker(guild.id, nonce=request.nonce) await self.chunker(guild.id, nonce=request.nonce)
await request.future if wait:
await request.future
return request.future
async def _chunk_and_dispatch(self, guild, unavailable): async def _chunk_and_dispatch(self, guild, unavailable):
try: try:
@ -1041,20 +1053,32 @@ class AutoShardedConnectionState(ConnectionState):
except asyncio.TimeoutError: except asyncio.TimeoutError:
break break
else: else:
try: if self._fetch_offline:
if self._fetch_offline: # Chunk the guild in the background while we wait for GUILD_CREATE streaming
await asyncio.wait_for(self.chunk_guild(guild), timeout=60.0) future = asyncio.ensure_future(self.chunk_guild(guild))
except asyncio.TimeoutError: else:
log.info('Timed out waiting for chunks while launching ready event.') future = self.loop.create_future()
finally: future.set_result(True)
processed.append(guild)
if guild.unavailable is False: processed.append((guild, future))
self.dispatch('guild_available', guild)
else: guilds = sorted(processed, key=lambda g: g[0].shard_id)
self.dispatch('guild_join', guild) for shard_id, info in itertools.groupby(guilds, key=lambda g: g[0].shard_id):
children, futures = zip(*info)
# 110 reqs/minute w/ 1 req/guild plus some buffer
timeout = 61 * (len(children) / 110)
try:
await utils.sane_wait_for(futures, timeout=timeout)
except asyncio.TimeoutError:
log.warning('Shard ID %s failed to wait for chunks (timeout=%.2f) for %d guilds', self.shard_id,
timeout,
len(guilds))
for guild in children:
if guild.unavailable is False:
self.dispatch('guild_available', guild)
else:
self.dispatch('guild_join', guild)
guilds = sorted(processed, key=lambda g: g.shard_id)
for shard_id, _ in itertools.groupby(guilds, key=lambda g: g.shard_id):
self.dispatch('shard_ready', shard_id) self.dispatch('shard_ready', shard_id)
# remove the state # remove the state