PYTHON-5044 - Successive AsyncMongoClients on a single loop always timeout on server selection

Author: NoahStapp

pymongo/asynchronous/mongo_client.py

@@ -1565,6 +1565,8 @@ async def close(self) -> None:
             # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
             await self._encrypter.close()
         self._closed = True
+        # Yield to the asyncio event loop so all executor tasks properly exit after cancellation
+        await asyncio.sleep(0)
 
     if not _IS_SYNC:
         # Add support for contextlib.aclosing.

pymongo/network_layer.py

@@ -267,18 +267,25 @@ async def async_receive_data(
         else:
             read_task = create_task(_async_receive(sock, length, loop))  # type: ignore[arg-type]
         tasks = [read_task, cancellation_task]
-        done, pending = await asyncio.wait(
-            tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
-        )
-        for task in pending:
-            task.cancel()
-        if pending:
-            await asyncio.wait(pending)
-        if len(done) == 0:
-            raise socket.timeout("timed out")
-        if read_task in done:
-            return read_task.result()
-        raise _OperationCancelled("operation cancelled")
+        try:
+            done, pending = await asyncio.wait(
+                tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED
+            )
+            for task in pending:
+                task.cancel()
+            if pending:
+                await asyncio.wait(pending)
+            if len(done) == 0:
+                raise socket.timeout("timed out")
+            if read_task in done:
+                return read_task.result()
+            raise _OperationCancelled("operation cancelled")
+        except asyncio.CancelledError:
+            for task in tasks:
+                task.cancel()
+            await asyncio.wait(tasks)
+            raise
+
     finally:
         sock.settimeout(sock_timeout)
 

pymongo/periodic_executor.py

@@ -75,6 +75,8 @@ def close(self, dummy: Any = None) -> None:
         callback; see monitor.py.
         """
         self._stopped = True
+        if self._task:
+            self._task.cancel()
 
     async def join(self, timeout: Optional[int] = None) -> None:
         if self._task is not None:

pymongo/synchronous/mongo_client.py

@@ -1559,6 +1559,7 @@ def close(self) -> None:
             # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
             self._encrypter.close()
         self._closed = True
+        # Yield to the asyncio event loop so all executor tasks properly exit after cancellation
 
     if not _IS_SYNC:
         # Add support for contextlib.closing.