PYTHON-5053 - AsyncMongoClient.close() should await all background tasks

pymongo/asynchronous/mongo_client.py

@@ -1565,6 +1565,15 @@ async def close(self) -> None:
             # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
             await self._encrypter.close()
         self._closed = True
+        if not _IS_SYNC:
+            join_tasks = [self._kill_cursors_executor]
+            try:
+                while self._topology._monitor_tasks:
+                    join_tasks.append(self._topology._monitor_tasks.pop())
+            except IndexError:
+                pass
+            join_tasks = [t.join() for t in join_tasks]  # type: ignore[func-returns-value]
+            await asyncio.gather(*join_tasks)
 
     if not _IS_SYNC:
         # Add support for contextlib.aclosing.

pymongo/asynchronous/monitor.py

@@ -112,9 +112,9 @@ async def close(self) -> None:
         """
         self.gc_safe_close()
 
-    async def join(self, timeout: Optional[int] = None) -> None:
+    async def join(self) -> None:
         """Wait for the monitor to stop."""
-        await self._executor.join(timeout)
+        await self._executor.join()
 
     def request_check(self) -> None:
         """If the monitor is sleeping, wake it soon."""
@@ -189,6 +189,9 @@ def gc_safe_close(self) -> None:
         self._rtt_monitor.gc_safe_close()
         self.cancel_check()
 
+    async def join(self) -> None:
+        await asyncio.gather(*[self._executor.join(), self._rtt_monitor.join()])
+
     async def close(self) -> None:
         self.gc_safe_close()
         await self._rtt_monitor.close()

pymongo/asynchronous/topology.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 import queue
@@ -29,7 +30,7 @@
 
 from pymongo import _csot, common, helpers_shared, periodic_executor
 from pymongo.asynchronous.client_session import _ServerSession, _ServerSessionPool
-from pymongo.asynchronous.monitor import SrvMonitor
+from pymongo.asynchronous.monitor import MonitorBase, SrvMonitor
 from pymongo.asynchronous.pool import Pool
 from pymongo.asynchronous.server import Server
 from pymongo.errors import (
@@ -207,6 +208,9 @@ async def target() -> bool:
         if self._settings.fqdn is not None and not self._settings.load_balanced:
             self._srv_monitor = SrvMonitor(self, self._settings)
 
+        # Stores all monitor tasks that need to be joined on close or server selection
+        self._monitor_tasks: list[MonitorBase] = []
+
     async def open(self) -> None:
         """Start monitoring, or restart after a fork.
 
@@ -241,6 +245,8 @@ async def open(self) -> None:
                 # Close servers and clear the pools.
                 for server in self._servers.values():
                     await server.close()
+                    if not _IS_SYNC:
+                        self._monitor_tasks.append(server._monitor)
                 # Reset the session pool to avoid duplicate sessions in
                 # the child process.
                 self._session_pool.reset()
@@ -288,10 +294,22 @@ async def select_servers(
                 selector, server_timeout, operation, operation_id, address
             )
 
-            return [
+            servers = [
                 cast(Server, self.get_server_by_address(sd.address)) for sd in server_descriptions
             ]
 
+        if not _IS_SYNC and self._monitor_tasks:
+            join_tasks = []
+            try:
+                while self._monitor_tasks:
+                    join_tasks.append(self._monitor_tasks.pop())
+            except IndexError:
+                pass
+            join_tasks = [t.join() for t in join_tasks]  # type: ignore[func-returns-value]
+            await asyncio.gather(*join_tasks)
+
+        return servers
+
     async def _select_servers_loop(
         self,
         selector: Callable[[Selection], Selection],
@@ -520,6 +538,8 @@ async def _process_change(
             and self._description.topology_type not in SRV_POLLING_TOPOLOGIES
         ):
             await self._srv_monitor.close()
+            if not _IS_SYNC:
+                self._monitor_tasks.append(self._srv_monitor)
 
         # Clear the pool from a failed heartbeat.
         if reset_pool:
@@ -695,6 +715,8 @@ async def close(self) -> None:
             old_td = self._description
             for server in self._servers.values():
                 await server.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(server._monitor)
 
             # Mark all servers Unknown.
             self._description = self._description.reset()
@@ -705,6 +727,8 @@ async def close(self) -> None:
             # Stop SRV polling thread.
             if self._srv_monitor:
                 await self._srv_monitor.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(self._srv_monitor)
 
             self._opened = False
             self._closed = True
@@ -944,6 +968,8 @@ async def _update_servers(self) -> None:
         for address, server in list(self._servers.items()):
             if not self._description.has_server(address):
                 await server.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(server._monitor)
                 self._servers.pop(address)
 
     def _create_pool_for_server(self, address: _Address) -> Pool:

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 is not None:
+            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,15 @@ def close(self) -> None:
             # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
             self._encrypter.close()
         self._closed = True
+        if not _IS_SYNC:
+            join_tasks = [self._kill_cursors_executor]
+            try:
+                while self._topology._monitor_tasks:
+                    join_tasks.append(self._topology._monitor_tasks.pop())
+            except IndexError:
+                pass
+            join_tasks = [t.join() for t in join_tasks]  # type: ignore[func-returns-value]
+            asyncio.gather(*join_tasks)
 
     if not _IS_SYNC:
         # Add support for contextlib.closing.

pymongo/synchronous/monitor.py

@@ -112,9 +112,9 @@ def close(self) -> None:
         """
         self.gc_safe_close()
 
-    def join(self, timeout: Optional[int] = None) -> None:
+    def join(self) -> None:
         """Wait for the monitor to stop."""
-        self._executor.join(timeout)
+        self._executor.join()
 
     def request_check(self) -> None:
         """If the monitor is sleeping, wake it soon."""
@@ -189,6 +189,9 @@ def gc_safe_close(self) -> None:
         self._rtt_monitor.gc_safe_close()
         self.cancel_check()
 
+    def join(self) -> None:
+        asyncio.gather(*[self._executor.join(), self._rtt_monitor.join()])
+
     def close(self) -> None:
         self.gc_safe_close()
         self._rtt_monitor.close()

pymongo/synchronous/topology.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 import queue
@@ -61,7 +62,7 @@
     writable_server_selector,
 )
 from pymongo.synchronous.client_session import _ServerSession, _ServerSessionPool
-from pymongo.synchronous.monitor import SrvMonitor
+from pymongo.synchronous.monitor import MonitorBase, SrvMonitor
 from pymongo.synchronous.pool import Pool
 from pymongo.synchronous.server import Server
 from pymongo.topology_description import (
@@ -207,6 +208,9 @@ def target() -> bool:
         if self._settings.fqdn is not None and not self._settings.load_balanced:
             self._srv_monitor = SrvMonitor(self, self._settings)
 
+        # Stores all monitor tasks that need to be joined on close or server selection
+        self._monitor_tasks: list[MonitorBase] = []
+
     def open(self) -> None:
         """Start monitoring, or restart after a fork.
 
@@ -241,6 +245,8 @@ def open(self) -> None:
                 # Close servers and clear the pools.
                 for server in self._servers.values():
                     server.close()
+                    if not _IS_SYNC:
+                        self._monitor_tasks.append(server._monitor)
                 # Reset the session pool to avoid duplicate sessions in
                 # the child process.
                 self._session_pool.reset()
@@ -288,10 +294,22 @@ def select_servers(
                 selector, server_timeout, operation, operation_id, address
             )
 
-            return [
+            servers = [
                 cast(Server, self.get_server_by_address(sd.address)) for sd in server_descriptions
             ]
 
+        if not _IS_SYNC and self._monitor_tasks:
+            join_tasks = []
+            try:
+                while self._monitor_tasks:
+                    join_tasks.append(self._monitor_tasks.pop())
+            except IndexError:
+                pass
+            join_tasks = [t.join() for t in join_tasks]  # type: ignore[func-returns-value]
+            asyncio.gather(*join_tasks)
+
+        return servers
+
     def _select_servers_loop(
         self,
         selector: Callable[[Selection], Selection],
@@ -520,6 +538,8 @@ def _process_change(
             and self._description.topology_type not in SRV_POLLING_TOPOLOGIES
         ):
             self._srv_monitor.close()
+            if not _IS_SYNC:
+                self._monitor_tasks.append(self._srv_monitor)
 
         # Clear the pool from a failed heartbeat.
         if reset_pool:
@@ -693,6 +713,8 @@ def close(self) -> None:
             old_td = self._description
             for server in self._servers.values():
                 server.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(server._monitor)
 
             # Mark all servers Unknown.
             self._description = self._description.reset()
@@ -703,6 +725,8 @@ def close(self) -> None:
             # Stop SRV polling thread.
             if self._srv_monitor:
                 self._srv_monitor.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(self._srv_monitor)
 
             self._opened = False
             self._closed = True
@@ -942,6 +966,8 @@ def _update_servers(self) -> None:
         for address, server in list(self._servers.items()):
             if not self._description.has_server(address):
                 server.close()
+                if not _IS_SYNC:
+                    self._monitor_tasks.append(server._monitor)
                 self._servers.pop(address)
 
     def _create_pool_for_server(self, address: _Address) -> Pool: