patch bug where Client.close() doesn't complete fully when called by user-code

discord/client.py

@@ -287,7 +287,7 @@ def __init__(self, *, intents: Intents, **options: Any) -> None:
         self._enable_debug_events: bool = options.pop('enable_debug_events', False)
         self._connection: ConnectionState[Self] = self._get_state(intents=intents, **options)
         self._connection.shard_count = self.shard_count
-        self._closed: bool = False
+        self._closure: Optional[asyncio.Task[None]] = None
         self._ready: asyncio.Event = MISSING
         self._application: Optional[AppInfo] = None
         self._connection._get_websocket = self._get_websocket
@@ -307,8 +307,7 @@ async def __aexit__(
         exc_value: Optional[BaseException],
         traceback: Optional[TracebackType],
     ) -> None:
-        if not self.is_closed():
-            await self.close()
+        await self.close()
 
     # internals
 
@@ -653,7 +652,7 @@ async def connect(self, *, reconnect: bool = True) -> None:
             'initial': True,
             'shard_id': self.shard_id,
         }
-        while not self.is_closed():
+        while not self._closure:
             try:
                 coro = DiscordWebSocket.from_client(self, **ws_params)
                 self.ws = await asyncio.wait_for(coro, timeout=60.0)
@@ -683,7 +682,7 @@ async def connect(self, *, reconnect: bool = True) -> None:
                         return
                     raise
 
-                if self.is_closed():
+                if self._closure is not None:
                     return
 
                 # If we get connection reset by peer then try to RESUME
@@ -721,16 +720,7 @@ async def connect(self, *, reconnect: bool = True) -> None:
                     session=self.ws.session_id,
                 )
 
-    async def close(self) -> None:
-        """|coro|
-
-        Closes the connection to Discord.
-        """
-        if self._closed:
-            return
-
-        self._closed = True
-
+    async def _close(self) -> None:
         await self._connection.close()
 
         if self.ws is not None and self.ws.open:
@@ -743,14 +733,24 @@ async def close(self) -> None:
 
         self.loop = MISSING
 
+    async def close(self) -> None:
+        """|coro|
+
+        Closes the connection to Discord.
+        """
+        if self._closure is None:
+            self._closure = self.loop.create_task(self._close())
+
+        await self._closure
+
     def clear(self) -> None:
         """Clears the internal state of the bot.
 
         After this, the bot can be considered "re-opened", i.e. :meth:`is_closed`
         and :meth:`is_ready` both return ``False`` along with the bot's internal
         cache cleared.
         """
-        self._closed = False
+        self._closure = None
         self._ready.clear()
         self._connection.clear()
         self.http.clear()
@@ -870,7 +870,7 @@ async def runner():
 
     def is_closed(self) -> bool:
         """:class:`bool`: Indicates if the websocket connection is closed."""
-        return self._closed
+        return self._closure is not None and self._closure.done()
 
     @property
     def activity(self) -> Optional[ActivityTypes]:

discord/shard.py

@@ -429,7 +429,7 @@ async def launch_shard(self, gateway: yarl.URL, shard_id: int, *, initial: bool
         ret.launch()
 
     async def launch_shards(self) -> None:
-        if self.is_closed():
+        if self._closure is not None:
             return
 
         if self.shard_count is None:
@@ -456,7 +456,7 @@ async def connect(self, *, reconnect: bool = True) -> None:
         self._reconnect = reconnect
         await self.launch_shards()
 
-        while not self.is_closed():
+        while not self._closure:
             item = await self.__queue.get()
             if item.type == EventType.close:
                 await self.close()
@@ -476,15 +476,7 @@ async def connect(self, *, reconnect: bool = True) -> None:
             elif item.type == EventType.clean_close:
                 return
 
-    async def close(self) -> None:
-        """|coro|
-
-        Closes the connection to Discord.
-        """
-        if self.is_closed():
-            return
-
-        self._closed = True
+    async def _close(self) -> None:
         await self._connection.close()
 
         to_close = [asyncio.ensure_future(shard.close(), loop=self.loop) for shard in self.__shards.values()]