Skip to content

Commit

Permalink
Update message references in AutoShardedConnectionState
Browse files Browse the repository at this point in the history
  • Loading branch information
Rapptz committed Jul 25, 2020
1 parent bf02831 commit 7d4e36f
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 0 deletions.
8 changes: 8 additions & 0 deletions discord/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -524,6 +524,14 @@ def _handle_call(self, call):
call['participants'] = participants
self.call = CallMessage(message=self, **call)

def _rebind_channel_reference(self, new_channel):
self.channel = new_channel

try:
del self._cs_guild
except AttributeError:
pass

@utils.cached_slot_property('_cs_guild')
def guild(self):
"""Optional[:class:`Guild`]: The guild that the message belongs to, if applicable."""
Expand Down
23 changes: 23 additions & 0 deletions discord/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -1062,6 +1062,17 @@ def __init__(self, *args, **kwargs):
self.shard_ids = ()
self.shards_launched = asyncio.Event()

def _update_message_references(self):
for msg in self._messages:
if not msg.guild:
continue

new_guild = self._get_guild(msg.guild.id)
if new_guild is not None and new_guild is not msg.guild:
channel_id = msg.channel.id
channel = new_guild.get_channel(channel_id) or Object(id=channel_id)
msg._rebind_channel_reference(channel)

async def chunker(self, guild_id, query='', limit=0, *, shard_id=None, nonce=None):
ws = self._get_websocket(guild_id, shard_id=shard_id)
await ws.request_chunks(guild_id, query=query, limit=limit, nonce=nonce)
Expand Down Expand Up @@ -1141,12 +1152,24 @@ def parse_ready(self, data):
if guild.large:
guilds.append((guild, guild.unavailable))

if self._messages:
self._update_message_references()

for pm in data.get('private_channels', []):
factory, _ = _channel_factory(pm['type'])
self._add_private_channel(factory(me=user, data=pm, state=self))

self.dispatch('connect')
self.dispatch('shard_connect', data['__shard_id__'])

# Much like clear(), if we have a massive deallocation
# then it's better to explicitly call the GC
# Note that in the original ready parsing code this was done
# implicitly via clear() but in the auto sharded client clearing
# the cache would have the consequence of clearing data on other
# shards as well.
gc.collect()

if self._ready_task is None:
self._ready_task = asyncio.ensure_future(self._delay_ready(), loop=self.loop)

Expand Down

0 comments on commit 7d4e36f

Please sign in to comment.