PYTHON-5044 - Fix successive AsyncMongoClients on a single loop always ti…

pymongo/asynchronous/mongo_client.py

@@ -1559,6 +1559,8 @@ async def close(self) -> None:
         # Stop the periodic task thread and then send pending killCursor
         # requests before closing the topology.
         self._kill_cursors_executor.close()
+        if not _IS_SYNC:
+            await self._kill_cursors_executor.join()
         await self._process_kill_cursors()
         await self._topology.close()
         if self._encrypter:

pymongo/asynchronous/monitor.py

@@ -191,6 +191,8 @@ def gc_safe_close(self) -> None:
 
     async def close(self) -> None:
         self.gc_safe_close()
+        if not _IS_SYNC:
+            await self._executor.join()
         await self._rtt_monitor.close()
         # Increment the generation and maybe close the socket. If the executor
         # thread has the socket checked out, it will be closed when checked in.
@@ -458,6 +460,8 @@ def __init__(self, topology: Topology, topology_settings: TopologySettings, pool
 
     async def close(self) -> None:
         self.gc_safe_close()
+        if not _IS_SYNC:
+            await self._executor.join()
         # Increment the generation and maybe close the socket. If the executor
         # thread has the socket checked out, it will be closed when checked in.
         await self._pool.reset()

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,17 +75,12 @@ 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:
-            try:
-                await asyncio.wait_for(self._task, timeout=timeout)  # type-ignore: [arg-type]
-            except asyncio.TimeoutError:
-                # Task timed out
-                pass
-            except asyncio.exceptions.CancelledError:
-                # Task was already finished, or not yet started.
-                raise
+            await asyncio.wait([self._task], timeout=timeout)  # type-ignore: [arg-type]
 
     def wake(self) -> None:
         """Execute the target function soon."""

pymongo/synchronous/mongo_client.py

@@ -1553,6 +1553,8 @@ def close(self) -> None:
         # Stop the periodic task thread and then send pending killCursor
         # requests before closing the topology.
         self._kill_cursors_executor.close()
+        if not _IS_SYNC:
+            self._kill_cursors_executor.join()
         self._process_kill_cursors()
         self._topology.close()
         if self._encrypter:

pymongo/synchronous/monitor.py

@@ -191,6 +191,8 @@ def gc_safe_close(self) -> None:
 
     def close(self) -> None:
         self.gc_safe_close()
+        if not _IS_SYNC:
+            self._executor.join()
         self._rtt_monitor.close()
         # Increment the generation and maybe close the socket. If the executor
         # thread has the socket checked out, it will be closed when checked in.
@@ -458,6 +460,8 @@ def __init__(self, topology: Topology, topology_settings: TopologySettings, pool
 
     def close(self) -> None:
         self.gc_safe_close()
+        if not _IS_SYNC:
+            self._executor.join()
         # Increment the generation and maybe close the socket. If the executor
         # thread has the socket checked out, it will be closed when checked in.
         self._pool.reset()

test/__init__.py

@@ -864,16 +864,6 @@ def max_message_size_bytes(self):
 client_context = ClientContext()
 
 
-def reset_client_context():
-    if _IS_SYNC:
-        # sync tests don't need to reset a client context
-        return
-    elif client_context.client is not None:
-        client_context.client.close()
-        client_context.client = None
-    client_context._init_client()
-
-
 class PyMongoTestCase(unittest.TestCase):
     def assertEqualCommand(self, expected, actual, msg=None):
         self.assertEqual(sanitize_cmd(expected), sanitize_cmd(actual), msg)
@@ -1136,8 +1126,6 @@ class IntegrationTest(PyMongoTestCase):
 
     @client_context.require_connection
     def setUp(self) -> None:
-        if not _IS_SYNC:
-            reset_client_context()
         if client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
         if client_context.serverless and not getattr(self, "RUN_ON_SERVERLESS", False):

test/asynchronous/__init__.py

@@ -866,16 +866,6 @@ async def max_message_size_bytes(self):
 async_client_context = AsyncClientContext()
 
 
-async def reset_client_context():
-    if _IS_SYNC:
-        # sync tests don't need to reset a client context
-        return
-    elif async_client_context.client is not None:
-        await async_client_context.client.close()
-        async_client_context.client = None
-    await async_client_context._init_client()
-
-
 class AsyncPyMongoTestCase(unittest.IsolatedAsyncioTestCase):
     def assertEqualCommand(self, expected, actual, msg=None):
         self.assertEqual(sanitize_cmd(expected), sanitize_cmd(actual), msg)
@@ -1154,8 +1144,6 @@ class AsyncIntegrationTest(AsyncPyMongoTestCase):
 
     @async_client_context.require_connection
     async def asyncSetUp(self) -> None:
-        if not _IS_SYNC:
-            await reset_client_context()
         if async_client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
         if async_client_context.serverless and not getattr(self, "RUN_ON_SERVERLESS", False):

test/asynchronous/test_connections_survive_primary_stepdown_spec.py

@@ -22,7 +22,6 @@
 from test.asynchronous import (
     AsyncIntegrationTest,
     async_client_context,
-    reset_client_context,
     unittest,
 )
 from test.asynchronous.helpers import async_repl_set_step_down

test/asynchronous/utils_spec_runner.py

@@ -647,10 +647,11 @@ async def setup_scenario(self, scenario_def):
     async def run_scenario(self, scenario_def, test):
         self.maybe_skip_scenario(test)
 
-        # Kill all sessions before and after each test to prevent an open
+        # Kill all sessions before (sync only) and after each test to prevent an open
         # transaction (from a test failure) from blocking collection/database
         # operations during test set up and tear down.
-        await self.kill_all_sessions()
+        if _IS_SYNC:
+            await self.kill_all_sessions()
         self.addAsyncCleanup(self.kill_all_sessions)
         await self.setup_scenario(scenario_def)
         database_name = self.get_scenario_db_name(scenario_def)

test/test_connections_survive_primary_stepdown_spec.py

@@ -22,7 +22,6 @@
 from test import (
     IntegrationTest,
     client_context,
-    reset_client_context,
     unittest,
 )
 from test.helpers import repl_set_step_down

test/utils_spec_runner.py

@@ -647,10 +647,11 @@ def setup_scenario(self, scenario_def):
     def run_scenario(self, scenario_def, test):
         self.maybe_skip_scenario(test)
 
-        # Kill all sessions before and after each test to prevent an open
+        # Kill all sessions before (sync only) and after each test to prevent an open
         # transaction (from a test failure) from blocking collection/database
         # operations during test set up and tear down.
-        self.kill_all_sessions()
+        if _IS_SYNC:
+            self.kill_all_sessions()
         self.addCleanup(self.kill_all_sessions)
         self.setup_scenario(scenario_def)
         database_name = self.get_scenario_db_name(scenario_def)