From 0401ab561b8857f532b8ef6dfe96fcef4450b7e6 Mon Sep 17 00:00:00 2001
From: Rapptz <rapptz@gmail.com>
Date: Sun, 10 Apr 2016 00:47:00 -0400
Subject: [PATCH] Proper chunking for unavailable guilds.
This will also delay on_ready until all chunking is complete.
---
discord/state.py | 90 +++++++++++++++++++++++++++++++++++-------------
1 file changed, 66 insertions(+), 24 deletions(-)
diff --git a/discord/state.py b/discord/state.py
index fa296f823..4baab8b68 100644
--- a/discord/state.py
+++ b/discord/state.py
@@ -46,6 +46,7 @@ class ListenerType(enum.Enum):
Listener = namedtuple('Listener', ('type', 'future', 'predicate'))
log = logging.getLogger(__name__)
+ReadyState = namedtuple('ReadyState', ('launch', 'chunks'))
class ConnectionState:
def __init__(self, dispatch, chunker, max_messages, *, loop):
@@ -133,33 +134,50 @@ class ConnectionState:
yield self.receive_chunk(server.id)
@asyncio.coroutine
- def _fill_offline(self):
- # a chunk has a maximum of 1000 members.
- # we need to find out how many futures we're actually waiting for
- large_servers = [s for s in self.servers if s.large]
- yield from self.chunker(large_servers)
-
- chunks = []
- for server in large_servers:
- chunks.extend(self.chunks_needed(server))
-
+ def _delay_ready(self, large_servers):
+ if len(large_servers):
+ # for regular accounts with < 100 guilds you will
+ # get a regular READY packet without the unavailable
+ # streaming so if this is non-empty then it's a regular
+ # account and it needs chunking.
+ yield from self.chunker(large_servers)
+ for server in large_servers:
+ self._ready_state.chunks.extend(self.chunks_needed(server))
+
+ launch = self._ready_state.launch
+ while not launch.is_set():
+ # this snippet of code is basically waiting 2 seconds
+ # until the last GUILD_CREATE was sent
+ launch.set()
+ yield from asyncio.sleep(2)
+
+ # get all the chunks
+ chunks = [f for f in self._ready_state.chunks if not f.done()]
if chunks:
yield from asyncio.wait(chunks)
+ # remove the state
+ del self._ready_state
+
+ # dispatch the event
self.dispatch('ready')
def parse_ready(self, data):
+ self._ready_state = ReadyState(launch=asyncio.Event(), chunks=[])
self.user = User(**data['user'])
guilds = data.get('guilds')
+ large_servers = []
for guild in guilds:
- self._add_server_from_data(guild)
+ server = self._add_server_from_data(guild)
+ if server.large:
+ large_servers.append(server)
for pm in data.get('private_channels'):
self._add_private_channel(PrivateChannel(id=pm['id'],
user=User(**pm['recipient'])))
- utils.create_task(self._fill_offline(), loop=self.loop)
+ utils.create_task(self._delay_ready(large_servers), loop=self.loop)
def parse_message_create(self, data):
channel = self.get_channel(data.get('channel_id'))
@@ -295,10 +313,8 @@ class ConnectionState:
self.dispatch('member_update', old_member, member)
- @asyncio.coroutine
- def parse_guild_create(self, data):
- unavailable = data.get('unavailable')
- if unavailable == False:
+ def _get_create_server(self, data):
+ if data.get('unavailable') == False:
# GUILD_CREATE with unavailable in the response
# usually means that the server has become available
# and is therefore in the cache
@@ -306,27 +322,53 @@ class ConnectionState:
if server is not None:
server.unavailable = False
server._from_data(data)
- self.dispatch('server_available', server)
- return
+ return server
+
+ return self._add_server_from_data(data)
+ @asyncio.coroutine
+ def parse_guild_create(self, data):
+ unavailable = data.get('unavailable')
if unavailable == True:
# joined a server with unavailable == True so..
return
- # if we're at this point then it was probably
- # unavailable during the READY event and is now
- # available, so it isn't in the cache...
-
- server = self._add_server_from_data(data)
+ server = self._get_create_server(data)
# check if it requires chunking
if server.large:
yield from self.chunker(server)
chunks = list(self.chunks_needed(server))
+
+ if unavailable == False:
+ # check if we're waiting for 'useful' READY
+ # and if we are, we don't want to dispatch any
+ # event such as server_join or server_available
+ # because we're still in the 'READY' phase. Or
+ # so we say.
+ try:
+ state = self._ready_state
+ state.launch.clear()
+ if chunks:
+ state.chunks.extend(chunks)
+ except AttributeError:
+ # the _ready_state attribute is only there during
+ # processing of useful READY.
+ pass
+ else:
+ return
+
+ # since we're not waiting for 'useful' READY we'll just
+ # do the chunk request here
if chunks:
yield from asyncio.wait(chunks)
- self.dispatch('server_join', server)
+
+ # Dispatch available if newly available
+ if unavailable == False:
+ self.dispatch('server_available', server)
+ else:
+ self.dispatch('server_join', server)
def parse_guild_update(self, data):
server = self._get_server(data.get('id'))