GH-111693: Propagate correct asyncio.CancelledError instance out of asyncio.Condition.wait()

Lib/asyncio/futures.py

@@ -138,9 +138,6 @@ def _make_cancelled_error(self):
             exc = exceptions.CancelledError()
         else:
             exc = exceptions.CancelledError(self._cancel_message)
-        exc.__context__ = self._cancelled_exc
-        # Remove the reference since we don't need this anymore.
-        self._cancelled_exc = None
         return exc
 
     def cancel(self, msg=None):

Lib/asyncio/locks.py

@@ -95,6 +95,9 @@ async def acquire(self):
         This method blocks until the lock is unlocked, then sets it to
         locked and returns True.
         """
+        # Implement fair scheduling, where thread always waits
+        # its turn.
+        # Jumping the queue if all are cancelled is an optimization.
         if (not self._locked and (self._waiters is None or
                 all(w.cancelled() for w in self._waiters))):
             self._locked = True
@@ -105,19 +108,16 @@ async def acquire(self):
         fut = self._get_loop().create_future()
         self._waiters.append(fut)
 
-        # Finally block should be called before the CancelledError
-        # handling as we don't want CancelledError to call
-        # _wake_up_first() and attempt to wake up itself.
         try:
-            try:
-                await fut
-            finally:
-                self._waiters.remove(fut)
-        except exceptions.CancelledError:
+            await fut
+        except BaseException:
+            self._waiters.remove(fut)
             if not self._locked:
+                # Error occurred after release() was called, must re-do release.
                 self._wake_up_first()
             raise
 
+        self._waiters.remove(fut)
         self._locked = True
         return True
 
@@ -269,17 +269,22 @@ async def wait(self):
                 self._waiters.remove(fut)
 
         finally:
-            # Must reacquire lock even if wait is cancelled
-            cancelled = False
+            # Must reacquire 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:
-                    cancelled = True
+                except exceptions.CancelledError as e:
+                    err = e
 
-            if cancelled:
-                raise exceptions.CancelledError
+            if err:
+                try:
+                    raise err  # re-raise same exception instance
+                finally:
+                    err = None  # brake reference cycles
 
     async def wait_for(self, predicate):
         """Wait until a predicate becomes true.
@@ -378,20 +383,18 @@ async def acquire(self):
         fut = self._get_loop().create_future()
         self._waiters.append(fut)
 
-        # Finally block should be called before the CancelledError
-        # handling as we don't want CancelledError to call
-        # _wake_up_first() and attempt to wake up itself.
         try:
-            try:
-                await fut
-            finally:
-                self._waiters.remove(fut)
-        except exceptions.CancelledError:
-            if not fut.cancelled():
+            await fut
+        except BaseException:
+            self._waiters.remove(fut)
+            if fut.done() and not fut.cancelled():
+                # Error occurred after release() was called for us.  Must undo
+                # the bookkeeping done there and retry.
                 self._value += 1
                 self._wake_up_next()
             raise
 
+        self._waiters.remove(fut)
         if self._value > 0:
             self._wake_up_next()
         return True

Lib/test/test_asyncio/test_locks.py

@@ -758,6 +758,63 @@ async def test_timeout_in_block(self):
             with self.assertRaises(asyncio.TimeoutError):
                 await asyncio.wait_for(condition.wait(), timeout=0.5)
 
+    async def test_cancelled_error_wakeup(self):
+        """Test that a cancelled error, received when awaiting wakeup
+        will be re-raised un-modified.
+        """
+        wake = False
+        raised = None
+        cond = asyncio.Condition()
+
+        async def func():
+            nonlocal raised
+            async with cond:
+                with self.assertRaises(asyncio.CancelledError) as err:
+                    await cond.wait_for(lambda: wake)
+                raised = err.exception
+                raise raised
+
+        task = asyncio.create_task(func())
+        await asyncio.sleep(0)
+        # Task is waiting on the condition, cancel it there
+        task.cancel(msg="foo")
+        with self.assertRaises(asyncio.CancelledError) as err:
+            await task
+        self.assertEqual(err.exception.args, ("foo",))
+        # we should have got the _same_ exception instance as the one originally raised
+        self.assertIs(err.exception, raised)
+
+    async def test_cancelled_error_re_aquire(self):
+        """Test that a cancelled error, received when re-aquiring lock,
+        will be re-raised un-modified.
+        """
+        wake = False
+        raised = None
+        cond = asyncio.Condition()
+
+        async def func():
+            nonlocal raised
+            async with cond:
+                with self.assertRaises(asyncio.CancelledError) as err:
+                    await cond.wait_for(lambda: wake)
+                raised = err.exception
+                raise raised
+
+        task = asyncio.create_task(func())
+        await asyncio.sleep(0)
+        # Task is waiting on the condition
+        await cond.acquire()
+        wake = True
+        cond.notify()
+        await asyncio.sleep(0)
+        # task is now trying to re-acquire the lock, cancel it there
+        task.cancel(msg="foo")
+        cond.release()
+        with self.assertRaises(asyncio.CancelledError) as err:
+            await task
+        self.assertEqual(err.exception.args, ("foo",))
+        # we should have got the _same_ exception instance as the one originally raised
+        self.assertIs(err.exception, raised)
 
 class SemaphoreTests(unittest.IsolatedAsyncioTestCase):