Skip to content

Commit

Permalink
Clear outstanding chunk requests when receiving a READY event
Browse files Browse the repository at this point in the history
Ref: #9571
  • Loading branch information
Rapptz committed Dec 24, 2023
1 parent 2a59e02 commit d25b574
Showing 1 changed file with 18 additions and 2 deletions.
20 changes: 18 additions & 2 deletions discord/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,12 +110,14 @@ class ChunkRequest:
def __init__(
self,
guild_id: int,
shard_id: int,
loop: asyncio.AbstractEventLoop,
resolver: Callable[[int], Any],
*,
cache: bool = True,
) -> None:
self.guild_id: int = guild_id
self.shard_id: int = shard_id
self.resolver: Callable[[int], Any] = resolver
self.loop: asyncio.AbstractEventLoop = loop
self.cache: bool = cache
Expand Down Expand Up @@ -315,6 +317,16 @@ def process_chunk_requests(self, guild_id: int, nonce: Optional[str], members: L
for key in removed:
del self._chunk_requests[key]

def clear_chunk_requests(self, shard_id: int | None) -> None:
removed = []
for key, request in self._chunk_requests.items():
if shard_id is None or request.shard_id == shard_id:
request.done()
removed.append(key)

for key in removed:
del self._chunk_requests[key]

def call_handlers(self, key: str, *args: Any, **kwargs: Any) -> None:
try:
func = self.handlers[key]
Expand Down Expand Up @@ -535,7 +547,7 @@ async def query_members(
if ws is None:
raise RuntimeError('Somehow do not have a websocket for this guild_id')

request = ChunkRequest(guild.id, self.loop, self._get_guild, cache=cache)
request = ChunkRequest(guild.id, guild.shard_id, self.loop, self._get_guild, cache=cache)
self._chunk_requests[request.nonce] = request

try:
Expand Down Expand Up @@ -602,6 +614,7 @@ def parse_ready(self, data: gw.ReadyEvent) -> None:

self._ready_state: asyncio.Queue[Guild] = asyncio.Queue()
self.clear(views=False)
self.clear_chunk_requests(None)
self.user = user = ClientUser(state=self, data=data['user'])
self._users[user.id] = user # type: ignore

Expand Down Expand Up @@ -1204,7 +1217,9 @@ async def chunk_guild(
cache = cache or self.member_cache_flags.joined
request = self._chunk_requests.get(guild.id)
if request is None:
self._chunk_requests[guild.id] = request = ChunkRequest(guild.id, self.loop, self._get_guild, cache=cache)
self._chunk_requests[guild.id] = request = ChunkRequest(
guild.id, guild.shard_id, self.loop, self._get_guild, cache=cache
)
await self.chunker(guild.id, nonce=request.nonce)

if wait:
Expand Down Expand Up @@ -1751,6 +1766,7 @@ def parse_ready(self, data: gw.ReadyEvent) -> None:

if shard_id in self._ready_tasks:
self._ready_tasks[shard_id].cancel()
self.clear_chunk_requests(shard_id)

if shard_id not in self._ready_states:
self._ready_states[shard_id] = asyncio.Queue()
Expand Down

0 comments on commit d25b574

Please sign in to comment.