PYTHON-4782 Fix deadlock and blocking behavior in _ACondition.wait

Author: ShaneHarvey

pymongo/lock.py

@@ -14,11 +14,12 @@
 from __future__ import annotations
 
 import asyncio
+import collections
 import os
 import threading
 import time
 import weakref
-from typing import Any, Callable, Optional
+from typing import Any, Optional
 
 _HAS_REGISTER_AT_FORK = hasattr(os, "register_at_fork")
 
@@ -44,6 +45,8 @@ def _release_locks() -> None:
 
 
 class _ALock:
+    __slots__ = ("_lock",)
+
     def __init__(self, lock: threading.Lock) -> None:
         self._lock = lock
 
@@ -82,8 +85,11 @@ async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
 
 
 class _ACondition:
+    __slots__ = ("_condition", "_waiters")
+
     def __init__(self, condition: threading.Condition) -> None:
         self._condition = condition
+        self._waiters = collections.deque()
 
     async def acquire(self, blocking: bool = True, timeout: float = -1) -> bool:
         if timeout > 0:
@@ -99,30 +105,115 @@ async def acquire(self, blocking: bool = True, timeout: float = -1) -> bool:
             await asyncio.sleep(0)
 
     async def wait(self, timeout: Optional[float] = None) -> bool:
-        if timeout is not None:
-            tstart = time.monotonic()
-        while True:
-            notified = self._condition.wait(0.001)
-            if notified:
-                return True
-            if timeout is not None and (time.monotonic() - tstart) > timeout:
-                return False
-
-    async def wait_for(self, predicate: Callable, timeout: Optional[float] = None) -> bool:
-        if timeout is not None:
-            tstart = time.monotonic()
-        while True:
-            notified = self._condition.wait_for(predicate, 0.001)
-            if notified:
-                return True
-            if timeout is not None and (time.monotonic() - tstart) > timeout:
-                return False
-
-    def notify(self, n: int = 1) -> None:
-        self._condition.notify(n)
-
-    def notify_all(self) -> None:
-        self._condition.notify_all()
+        """Wait until notified.
+
+        If the calling task has not acquired the lock when this
+        method is called, a RuntimeError is raised.
+
+        This method releases the underlying lock, and then blocks
+        until it is awakened by a notify() or notify_all() call for
+        the same condition variable in another task.  Once
+        awakened, it re-acquires the lock and returns True.
+
+        This method may return spuriously,
+        which is why the caller should always
+        re-check the state and be prepared to wait() again.
+        """
+        loop = asyncio.get_running_loop()
+        fut = loop.create_future()
+        self._waiters.append((loop, fut))
+        self.release()
+        try:
+            try:
+                try:
+                    await asyncio.wait_for(fut, timeout)
+                    return True
+                except asyncio.TimeoutError:
+                    return False  # Return false on timeout for sync pool compat.
+            finally:
+                # Must re-acquire lock even if wait is cancelled.
+                # We only catch CancelledError here, since we don't want any
+                # other (fatal) errors with the future to cause us to spin.
+                err = None
+                while True:
+                    try:
+                        await self.acquire()
+                        break
+                    except asyncio.exceptions.CancelledError as e:
+                        err = e
+
+                self._waiters.remove((loop, fut))
+                if err is not None:
+                    try:
+                        raise err  # Re-raise most recent exception instance.
+                    finally:
+                        err = None  # Break reference cycles.
+        except BaseException:
+            # Any error raised out of here _may_ have occurred after this Task
+            # believed to have been successfully notified.
+            # Make sure to notify another Task instead.  This may result
+            # in a "spurious wakeup", which is allowed as part of the
+            # Condition Variable protocol.
+            self.notify(1)
+            raise
+
+    async def wait_for(self, predicate):
+        """Wait until a predicate becomes true.
+
+        The predicate should be a callable which result will be
+        interpreted as a boolean value.  The final predicate value is
+        the return value.
+        """
+        result = predicate()
+        while not result:
+            await self.wait()
+            result = predicate()
+        return result
+
+    def notify(self, n=1):
+        """By default, wake up one coroutine waiting on this condition, if any.
+        If the calling coroutine has not acquired the lock when this method
+        is called, a RuntimeError is raised.
+
+        This method wakes up at most n of the coroutines waiting for the
+        condition variable; it is a no-op if no coroutines are waiting.
+
+        Note: an awakened coroutine does not actually return from its
+        wait() call until it can reacquire the lock. Since notify() does
+        not release the lock, its caller should.
+        """
+        idx = 0
+        to_remove = []
+        for loop, fut in self._waiters:
+            if idx >= n:
+                break
+
+            if fut.done():
+                continue
+
+            def safe_set_result(fut):
+                if not fut.done():
+                    fut.set_result(False)
+
+            try:
+                loop.call_soon_threadsafe(safe_set_result, fut)
+            except RuntimeError:
+                # Loop was closed, ignore.
+                to_remove.append((loop, fut))
+                continue
+
+            idx += 1
+
+        for waiter in to_remove:
+            self._waiters.remove(waiter)
+
+    def notify_all(self):
+        """Wake up all threads waiting on this condition. This method acts
+        like notify(), but wakes up all waiting threads instead of one. If the
+        calling thread has not acquired the lock when this method is called,
+        a RuntimeError is raised.
+        """
+        self.notify(len(self._waiters))
 
     def release(self) -> None:
         self._condition.release()