asyncio: preserve Task.cancel() msg through fut_waiter and CancelledError paths

Author: prince8273

Doc/whatsnew/3.16.rst

@@ -123,6 +123,11 @@ asyncio
   which has been deprecated since Python 3.14.
   Use :func:`inspect.iscoroutinefunction` instead.
 
+* :meth:`asyncio.Task.cancel` now preserves its *msg* argument when a task is
+  cancelled while waiting on another future, so the message is retained when
+  the task later raises :exc:`asyncio.CancelledError`.
+  (Contributed by Prince Kumar in :gh:`NNNNN`.)
+
 functools
 ---------
 

Lib/asyncio/tasks.py

@@ -216,6 +216,7 @@ def cancel(self, msg=None):
                 # Leave self._fut_waiter; it may be a Task that
                 # catches and ignores the cancellation so we may have
                 # to cancel it again later.
+                self._cancel_message = msg
                 return True
         # It must be the case that self.__step is already scheduled.
         self._must_cancel = True
@@ -299,7 +300,7 @@ def __step_run_and_handle_result(self, exc):
         except exceptions.CancelledError as exc:
             # Save the original exception so we can chain it later.
             self._cancelled_exc = exc
-            super().cancel()  # I.e., Future.cancel(self).
+            super().cancel(msg=self._cancel_message)  # I.e., Future.cancel(self).
         except (KeyboardInterrupt, SystemExit) as exc:
             super().set_exception(exc)
             raise

Lib/test/test_asyncio/test_tasks.py

@@ -1861,6 +1861,25 @@ async def coro():
         self.assertIsNone(task._fut_waiter)
         self.assertTrue(fut.cancelled())
 
+    def test_task_cancel_waiter_future_with_message(self):
+        fut = self.new_future(self.loop)
+
+        async def coro():
+            await fut
+
+        task = self.new_task(self.loop, coro())
+        test_utils.run_briefly(self.loop)
+        self.assertIs(task._fut_waiter, fut)
+
+        task.cancel('my message')
+        test_utils.run_briefly(self.loop)
+        with self.assertRaises(asyncio.CancelledError) as cm:
+            self.loop.run_until_complete(task)
+        self.assertEqual(cm.exception.args, ('my message',))
+        self.assertEqual(task._cancel_message, 'my message')
+        self.assertIsNone(task._fut_waiter)
+        self.assertTrue(fut.cancelled())
+
     def test_task_set_methods(self):
         async def notmuch():
             return 'ko'

Modules/_asynciomodule.c

@@ -2612,6 +2612,8 @@ _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
         }
 
         if (is_true) {
+            Py_XINCREF(msg);
+            Py_XSETREF(self->task_cancel_msg, msg);
             Py_RETURN_TRUE;
         }
     }
@@ -3136,7 +3138,7 @@ task_step_impl(asyncio_state *state, TaskObj *task, PyObject *exc)
             /* transfer ownership */
             fut->fut_cancelled_exc = exc;
 
-            return future_cancel(state, fut, NULL);
+            return future_cancel(state, fut, fut->fut_cancel_msg);
         }
 
         /* Some other exception; pop it and call Task.set_exception() */