PYTHON-4725 Async client should use tasks for SDAM instead of threads

pymongo/asynchronous/mongo_client.py

@@ -909,7 +909,7 @@ async def target() -> bool:
             await AsyncMongoClient._process_periodic_tasks(client)
             return True
 
-        executor = periodic_executor.PeriodicExecutor(
+        executor = periodic_executor.AsyncPeriodicExecutor(
             interval=common.KILL_CURSOR_FREQUENCY,
             min_interval=common.MIN_HEARTBEAT_INTERVAL,
             target=target,

pymongo/asynchronous/monitor.py

@@ -76,7 +76,7 @@ async def target() -> bool:
             await monitor._run()  # type:ignore[attr-defined]
             return True
 
-        executor = periodic_executor.PeriodicExecutor(
+        executor = periodic_executor.AsyncPeriodicExecutor(
             interval=interval, min_interval=min_interval, target=target, name=name
         )
 
@@ -112,9 +112,9 @@ async def close(self) -> None:
         """
         self.gc_safe_close()
 
-    def join(self, timeout: Optional[int] = None) -> None:
+    async def join(self, timeout: Optional[int] = None) -> None:
         """Wait for the monitor to stop."""
-        self._executor.join(timeout)
+        await self._executor.join(timeout)
 
     def request_check(self) -> None:
         """If the monitor is sleeping, wake it soon."""
@@ -250,7 +250,7 @@ async def _check_server(self) -> ServerDescription:
             except (OperationFailure, NotPrimaryError) as exc:
                 # Update max cluster time even when hello fails.
                 details = cast(Mapping[str, Any], exc.details)
-                self._topology.receive_cluster_time(details.get("$clusterTime"))
+                await self._topology.receive_cluster_time(details.get("$clusterTime"))
                 raise
         except ReferenceError:
             raise
@@ -531,4 +531,5 @@ def _shutdown_resources() -> None:
         shutdown()
 
 
-atexit.register(_shutdown_resources)
+if _IS_SYNC:
+    atexit.register(_shutdown_resources)

pymongo/asynchronous/periodic_executor.py

@@ -28,6 +28,99 @@
 _IS_SYNC = False
 
 
+class AsyncPeriodicExecutor:
+    def __init__(
+        self,
+        interval: float,
+        min_interval: float,
+        target: Any,
+        name: Optional[str] = None,
+    ):
+        """Run a target function periodically on a background thread.
+
+        If the target's return value is false, the executor stops.
+
+        :param interval: Seconds between calls to `target`.
+        :param min_interval: Minimum seconds between calls if `wake` is
+            called very often.
+        :param target: A function.
+        :param name: A name to give the underlying thread.
+        """
+        # threading.Event and its internal condition variable are expensive
+        # in Python 2, see PYTHON-983. Use a boolean to know when to wake.
+        # The executor's design is constrained by several Python issues, see
+        # "periodic_executor.rst" in this repository.
+        self._event = False
+        self._interval = interval
+        self._min_interval = min_interval
+        self._target = target
+        self._stopped = False
+        self._task: Optional[asyncio.Task] = None
+        self._name = name
+        self._skip_sleep = False
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}(name={self._name}) object at 0x{id(self):x}>"
+
+    def open(self) -> None:
+        """Start. Multiple calls have no effect."""
+        self._stopped = False
+        started = self._task and not self._task.done()
+
+        if not started:
+            self._task = asyncio.get_event_loop().create_task(self._run(), name=self._name)  # type: ignore[func-returns-value]
+
+    def close(self, dummy: Any = None) -> None:
+        """Stop. To restart, call open().
+
+        The dummy parameter allows an executor's close method to be a weakref
+        callback; see monitor.py.
+        """
+        self._stopped = True
+
+    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.
+                pass
+
+    def wake(self) -> None:
+        """Execute the target function soon."""
+        self._event = True
+
+    def update_interval(self, new_interval: int) -> None:
+        self._interval = new_interval
+
+    def skip_sleep(self) -> None:
+        self._skip_sleep = True
+
+    async def _run(self) -> None:
+        while not self._stopped:
+            try:
+                if not await self._target():
+                    self._stopped = True
+                    break
+            except BaseException:
+                self._stopped = True
+                raise
+
+            if self._skip_sleep:
+                self._skip_sleep = False
+            else:
+                deadline = time.monotonic() + self._interval
+                while not self._stopped and time.monotonic() < deadline:
+                    await asyncio.sleep(self._min_interval)
+                    if self._event:
+                        break  # Early wake.
+
+            self._event = False
+
+
 class PeriodicExecutor:
     def __init__(
         self,
@@ -64,19 +157,6 @@ def __init__(
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}(name={self._name}) object at 0x{id(self):x}>"
 
-    def _run_async(self) -> None:
-        # The default asyncio loop implementation on Windows
-        # has issues with sharing sockets across loops (https://github.com/python/cpython/issues/122240)
-        # We explicitly use a different loop implementation here to prevent that issue
-        if sys.platform == "win32":
-            loop = asyncio.SelectorEventLoop()
-            try:
-                loop.run_until_complete(self._run())  # type: ignore[func-returns-value]
-            finally:
-                loop.close()
-        else:
-            asyncio.run(self._run())  # type: ignore[func-returns-value]
-
     def open(self) -> None:
         """Start. Multiple calls have no effect.
 
@@ -104,10 +184,7 @@ def open(self) -> None:
             pass
 
         if not started:
-            if _IS_SYNC:
-                thread = threading.Thread(target=self._run, name=self._name)
-            else:
-                thread = threading.Thread(target=self._run_async, name=self._name)
+            thread = threading.Thread(target=self._run, name=self._name)
             thread.daemon = True
             self._thread = weakref.proxy(thread)
             _register_executor(self)

pymongo/asynchronous/topology.py

@@ -185,7 +185,7 @@ def __init__(self, topology_settings: TopologySettings):
             async def target() -> bool:
                 return process_events_queue(weak)
 
-            executor = periodic_executor.PeriodicExecutor(
+            executor = periodic_executor.AsyncPeriodicExecutor(
                 interval=common.EVENTS_QUEUE_FREQUENCY,
                 min_interval=common.MIN_HEARTBEAT_INTERVAL,
                 target=target,
@@ -654,7 +654,15 @@ async def request_check_all(self, wait_time: int = 5) -> None:
         """Wake all monitors, wait for at least one to check its server."""
         async with self._lock:
             self._request_check_all()
+            # if _IS_SYNC:
             await self._condition.wait(wait_time)
+            # else:
+            #     try:
+            #         await asyncio.wait_for(
+            #             self._condition.wait(), wait_time
+            #         )  # type-ignore: [arg-type]
+            #     except asyncio.TimeoutError:
+            #         pass
 
     def data_bearing_servers(self) -> list[ServerDescription]:
         """Return a list of all data-bearing servers.
@@ -742,7 +750,7 @@ async def close(self) -> None:
         if self._publish_server or self._publish_tp:
             # Make sure the events executor thread is fully closed before publishing the remaining events
             self.__events_executor.close()
-            self.__events_executor.join(1)
+            await self.__events_executor.join(1)
             process_events_queue(weakref.ref(self._events))  # type: ignore[arg-type]
 
     @property

pymongo/synchronous/monitor.py

@@ -531,4 +531,5 @@ def _shutdown_resources() -> None:
         shutdown()
 
 
-atexit.register(_shutdown_resources)
+if _IS_SYNC:
+    atexit.register(_shutdown_resources)

pymongo/synchronous/periodic_executor.py

@@ -28,6 +28,99 @@
 _IS_SYNC = True
 
 
+class SyncPeriodicExecutor:
+    def __init__(
+        self,
+        interval: float,
+        min_interval: float,
+        target: Any,
+        name: Optional[str] = None,
+    ):
+        """Run a target function periodically on a background thread.
+
+        If the target's return value is false, the executor stops.
+
+        :param interval: Seconds between calls to `target`.
+        :param min_interval: Minimum seconds between calls if `wake` is
+            called very often.
+        :param target: A function.
+        :param name: A name to give the underlying thread.
+        """
+        # threading.Event and its internal condition variable are expensive
+        # in Python 2, see PYTHON-983. Use a boolean to know when to wake.
+        # The executor's design is constrained by several Python issues, see
+        # "periodic_executor.rst" in this repository.
+        self._event = False
+        self._interval = interval
+        self._min_interval = min_interval
+        self._target = target
+        self._stopped = False
+        self._task: Optional[asyncio.Task] = None
+        self._name = name
+        self._skip_sleep = False
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}(name={self._name}) object at 0x{id(self):x}>"
+
+    def open(self) -> None:
+        """Start. Multiple calls have no effect."""
+        self._stopped = False
+        started = self._task and not self._task.done()
+
+        if not started:
+            self._task = asyncio.get_event_loop().create_task(self._run(), name=self._name)  # type: ignore[func-returns-value]
+
+    def close(self, dummy: Any = None) -> None:
+        """Stop. To restart, call open().
+
+        The dummy parameter allows an executor's close method to be a weakref
+        callback; see monitor.py.
+        """
+        self._stopped = True
+
+    def join(self, timeout: Optional[int] = None) -> None:
+        if self._task is not None:
+            try:
+                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.
+                pass
+
+    def wake(self) -> None:
+        """Execute the target function soon."""
+        self._event = True
+
+    def update_interval(self, new_interval: int) -> None:
+        self._interval = new_interval
+
+    def skip_sleep(self) -> None:
+        self._skip_sleep = True
+
+    def _run(self) -> None:
+        while not self._stopped:
+            try:
+                if not self._target():
+                    self._stopped = True
+                    break
+            except BaseException:
+                self._stopped = True
+                raise
+
+            if self._skip_sleep:
+                self._skip_sleep = False
+            else:
+                deadline = time.monotonic() + self._interval
+                while not self._stopped and time.monotonic() < deadline:
+                    time.sleep(self._min_interval)
+                    if self._event:
+                        break  # Early wake.
+
+            self._event = False
+
+
 class PeriodicExecutor:
     def __init__(
         self,
@@ -64,19 +157,6 @@ def __init__(
     def __repr__(self) -> str:
         return f"<{self.__class__.__name__}(name={self._name}) object at 0x{id(self):x}>"
 
-    def _run_async(self) -> None:
-        # The default asyncio loop implementation on Windows
-        # has issues with sharing sockets across loops (https://github.com/python/cpython/issues/122240)
-        # We explicitly use a different loop implementation here to prevent that issue
-        if sys.platform == "win32":
-            loop = asyncio.SelectorEventLoop()
-            try:
-                loop.run_until_complete(self._run())  # type: ignore[func-returns-value]
-            finally:
-                loop.close()
-        else:
-            asyncio.run(self._run())  # type: ignore[func-returns-value]
-
     def open(self) -> None:
         """Start. Multiple calls have no effect.
 
@@ -104,10 +184,7 @@ def open(self) -> None:
             pass
 
         if not started:
-            if _IS_SYNC:
-                thread = threading.Thread(target=self._run, name=self._name)
-            else:
-                thread = threading.Thread(target=self._run_async, name=self._name)
+            thread = threading.Thread(target=self._run, name=self._name)
             thread.daemon = True
             self._thread = weakref.proxy(thread)
             _register_executor(self)

pymongo/synchronous/topology.py

@@ -652,7 +652,15 @@ def request_check_all(self, wait_time: int = 5) -> None:
         """Wake all monitors, wait for at least one to check its server."""
         with self._lock:
             self._request_check_all()
+            # if _IS_SYNC:
             self._condition.wait(wait_time)
+            # else:
+            #     try:
+            #         asyncio.wait_for(
+            #             self._condition.wait(), wait_time
+            #         )  # type-ignore: [arg-type]
+            #     except asyncio.TimeoutError:
+            #         pass
 
     def data_bearing_servers(self) -> list[ServerDescription]:
         """Return a list of all data-bearing servers.

test/__init__.py

@@ -859,6 +859,16 @@ def max_message_size_bytes(self):
 client_context = ClientContext()
 
 
+def recreate_client_context():
+    if _IS_SYNC:
+        # sync tests don't need to recreate a client context
+        return
+    global client_context
+    teardown()
+    client_context = ClientContext()
+    setup()
+
+
 class PyMongoTestCase(unittest.TestCase):
     def assertEqualCommand(self, expected, actual, msg=None):
         self.assertEqual(sanitize_cmd(expected), sanitize_cmd(actual), msg)