gh-112202: Ensure that condition.notify() succeeds even when racing with Task.cancel()

Doc/library/asyncio-sync.rst

@@ -216,7 +216,7 @@ Condition
 
    .. method:: notify(n=1)
 
-      Wake up at most *n* tasks (1 by default) waiting on this
+      Wake up *n* tasks (1 by default) waiting on this
       condition.  The method is no-op if no tasks are waiting.
 
       The lock must be acquired before this method is called and
@@ -257,12 +257,18 @@ Condition
       Once awakened, the Condition re-acquires its lock and this method
       returns ``True``.
 
+      Note that a task *may* return from this call without having been
+      explicity awoken, which is why the caller should always re-check the state
+      and be prepared to :meth:`wait` again. For this reason, you may
+      prefer to use :meth:`wait_for` instead.
+
    .. coroutinemethod:: wait_for(predicate)
 
       Wait until a predicate becomes *true*.
 
       The predicate must be a callable which result will be
-      interpreted as a boolean value.  The final value is the
+      interpreted as a boolean value.  The method will repeatedly
+      :meth:`wait` until the predicate evaluates to *true*. The final value is the
       return value.
 
 

Lib/asyncio/locks.py

@@ -257,43 +257,57 @@ async def wait(self):
         until it is awakened by a notify() or notify_all() call for
         the same condition variable in another coroutine.  Once
         awakened, it re-acquires the lock and returns True.
+
+        This method may return without having been explicitly awoken by
+        a notify() or notify_all(), which is why the caller should always
+        re-check the state and be prepared to wait() again.
         """
         if not self.locked():
             raise RuntimeError('cannot wait on un-acquired lock')
 
+        fut = self._get_loop().create_future()
         self.release()
         try:
-            fut = self._get_loop().create_future()
-            self._waiters.append(fut)
             try:
-                await fut
-                return True
-            finally:
-                self._waiters.remove(fut)
-
-        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 exceptions.CancelledError as e:
-                    err = e
-
-            if err:
+                self._waiters.append(fut)
                 try:
-                    raise err  # Re-raise most recent exception instance.
+                    await fut
+                    return True
                 finally:
-                    err = None  # Break reference cycles.
+                    self._waiters.remove(fut)
+
+            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 exceptions.CancelledError as e:
+                        err = e
+
+                if err:
+                    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
+        interpreted as a boolean value.  The method will repeatedly
+        wait() until it evaluates to true.  The final predicate value is
         the return value.
         """
         result = predicate()
@@ -307,7 +321,7 @@ def notify(self, n=1):
         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
+        This method wakes up 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
@@ -316,7 +330,9 @@ def notify(self, n=1):
         """
         if not self.locked():
             raise RuntimeError('cannot notify on un-acquired lock')
+        self._notify(n)
 
+    def _notify(self, n):
         idx = 0
         for fut in self._waiters:
             if idx >= n:

Lib/test/test_asyncio/test_locks.py

@@ -816,6 +816,98 @@ async def func():
         # originally raised.
         self.assertIs(err.exception, raised)
 
+    async def test_cancelled_wakeup(self):
+        # Test that a task cancelled at the "same" time as it is woken
+        # up as part of a Condition.notify() does not result in a lost wakeup.
+        # This test simulates a cancel while the target task is awaiting initial
+        # wakeup on the wakeup queue.
+        condition = asyncio.Condition()
+        state = 0
+        async def consumer():
+            nonlocal state
+            async with condition:
+                while True:
+                    await condition.wait_for(lambda: state != 0)
+                    if state < 0:
+                        return
+                    state -= 1
+
+        # create two consumers
+        c = [asyncio.create_task(consumer()) for _ in range(2)]
+        # wait for them to settle
+        await asyncio.sleep(0)
+        async with condition:
+            # produce one item and wake up one
+            state += 1
+            condition.notify(1)
+
+            # Cancel it while it is awaiting to be run.
+            # This cancellation could come from the outside
+            c[0].cancel()
+
+            # now wait for the item to be consumed
+            # if it doesn't means that our "notify" didn"t take hold.
+            # because it raced with a cancel()
+            try:
+                async with asyncio.timeout(0.01):
+                    await condition.wait_for(lambda: state == 0)
+            except TimeoutError:
+                pass
+            self.assertEqual(state, 0)
+
+            # clean up
+            state = -1
+            condition.notify_all()
+        await c[1]
+
+    async def test_cancelled_wakeup_relock(self):
+        # Test that a task cancelled at the "same" time as it is woken
+        # up as part of a Condition.notify() does not result in a lost wakeup.
+        # This test simulates a cancel while the target task is acquiring the lock
+        # again.
+        condition = asyncio.Condition()
+        state = 0
+        async def consumer():
+            nonlocal state
+            async with condition:
+                while True:
+                    await condition.wait_for(lambda: state != 0)
+                    if state < 0:
+                        return
+                    state -= 1
+
+        # create two consumers
+        c = [asyncio.create_task(consumer()) for _ in range(2)]
+        # wait for them to settle
+        await asyncio.sleep(0)
+        async with condition:
+            # produce one item and wake up one
+            state += 1
+            condition.notify(1)
+
+            # now we sleep for a bit.  This allows the target task to wake up and
+            # settle on re-aquiring the lock
+            await asyncio.sleep(0)
+
+            # Cancel it while awaiting the lock
+            # This cancel could come the outside.
+            c[0].cancel()
+
+            # now wait for the item to be consumed
+            # if it doesn't means that our "notify" didn"t take hold.
+            # because it raced with a cancel()
+            try:
+                async with asyncio.timeout(0.01):
+                    await condition.wait_for(lambda: state == 0)
+            except TimeoutError:
+                pass
+            self.assertEqual(state, 0)
+
+            # clean up
+            state = -1
+            condition.notify_all()
+        await c[1]
+
 class SemaphoreTests(unittest.IsolatedAsyncioTestCase):
 
     def test_initial_value_zero(self):

Misc/NEWS.d/next/Library/2024-01-12-09-35-07.gh-issue-112202.t_0V1m.rst

@@ -0,0 +1 @@
+Ensure that a :func:`asyncio.Condition.notify` call does not get lost if the notified ``Task`` is simultaneously cancelled or encounters any other error.