gh-116720: Fix corner cases of taskgroups (#117407)

This prevents external cancellations of a task group's parent task to
be dropped when an internal cancellation happens at the same time.
Also strengthen the semantics of uncancel() to clear self._must_cancel
when the cancellation count reaches zero.

Co-Authored-By: Tin Tvrtković <tinchester@gmail.com>
Co-Authored-By: Arthur Tacca

8 files changed, 183 insertions, 13 deletions

diff --git a/Doc/library/asyncio-task.rst b/Doc/library/asyncio-task.rst
index 3b10a0d..3d300c3 100644
--- a/Doc/library/asyncio-task.rst
+++ b/Doc/library/asyncio-task.rst
@@ -392,6 +392,27 @@ is also included in the exception group.
 The same special case is made for
 :exc:`KeyboardInterrupt` and :exc:`SystemExit` as in the previous paragraph.
 
+Task groups are careful not to mix up the internal cancellation used to
+"wake up" their :meth:`~object.__aexit__` with cancellation requests
+for the task in which they are running made by other parties.
+In particular, when one task group is syntactically nested in another,
+and both experience an exception in one of their child tasks simultaneously,
+the inner task group will process its exceptions, and then the outer task group
+will receive another cancellation and process its own exceptions.
+
+In the case where a task group is cancelled externally and also must
+raise an :exc:`ExceptionGroup`, it will call the parent task's
+:meth:`~asyncio.Task.cancel` method. This ensures that a
+:exc:`asyncio.CancelledError` will be raised at the next
+:keyword:`await`, so the cancellation is not lost.
+
+Task groups preserve the cancellation count
+reported by :meth:`asyncio.Task.cancelling`.
+
+.. versionchanged:: 3.13
+
+   Improved handling of simultaneous internal and external cancellations
+   and correct preservation of cancellation counts.
 
 Sleeping
 ========
@@ -1369,6 +1390,15 @@ Task Object
       catching :exc:`CancelledError`, it needs to call this method to remove
       the cancellation state.
 
+      When this method decrements the cancellation count to zero,
+      the method checks if a previous :meth:`cancel` call had arranged
+      for :exc:`CancelledError` to be thrown into the task.
+      If it hasn't been thrown yet, that arrangement will be
+      rescinded (by resetting the internal ``_must_cancel`` flag).
+
+   .. versionchanged:: 3.13
+      Changed to rescind pending cancellation requests upon reaching zero.
+
    .. method:: cancelling()
 
       Return the number of pending cancellation requests to this Task, i.e.,
diff --git a/Doc/whatsnew/3.13.rst b/Doc/whatsnew/3.13.rst
index 707dcaa..d971beb 100644
--- a/Doc/whatsnew/3.13.rst
+++ b/Doc/whatsnew/3.13.rst
@@ -196,13 +196,6 @@ Other Language Changes
 
   (Contributed by Sebastian Pipping in :gh:`115623`.)
 
-* When :func:`asyncio.TaskGroup.create_task` is called on an inactive
-  :class:`asyncio.TaskGroup`, the given coroutine will be closed (which
-  prevents a :exc:`RuntimeWarning` about the given coroutine being
-  never awaited).
-
-  (Contributed by Arthur Tacca and Jason Zhang in :gh:`115957`.)
-
 * The :func:`ssl.create_default_context` API now includes
   :data:`ssl.VERIFY_X509_PARTIAL_CHAIN` and :data:`ssl.VERIFY_X509_STRICT`
   in its default flags.
@@ -300,6 +293,33 @@ asyncio
   with the tasks being completed.
   (Contributed by Justin Arthur in :gh:`77714`.)
 
+* When :func:`asyncio.TaskGroup.create_task` is called on an inactive
+  :class:`asyncio.TaskGroup`, the given coroutine will be closed (which
+  prevents a :exc:`RuntimeWarning` about the given coroutine being
+  never awaited).
+  (Contributed by Arthur Tacca and Jason Zhang in :gh:`115957`.)
+
+* Improved behavior of :class:`asyncio.TaskGroup` when an external cancellation
+  collides with an internal cancellation. For example, when two task groups
+  are nested and both experience an exception in a child task simultaneously,
+  it was possible that the outer task group would hang, because its internal
+  cancellation was swallowed by the inner task group.
+
+  In the case where a task group is cancelled externally and also must
+  raise an :exc:`ExceptionGroup`, it will now call the parent task's
+  :meth:`~asyncio.Task.cancel` method.  This ensures that a
+  :exc:`asyncio.CancelledError` will be raised at the next
+  :keyword:`await`, so the cancellation is not lost.
+
+  An added benefit of these changes is that task groups now preserve the
+  cancellation count (:meth:`asyncio.Task.cancelling`).
+
+  In order to handle some corner cases, :meth:`asyncio.Task.uncancel` may now
+  reset the undocumented ``_must_cancel`` flag when the cancellation count
+  reaches zero.
+
+  (Inspired by an issue reported by Arthur Tacca in :gh:`116720`.)
+
 * Add :meth:`asyncio.Queue.shutdown` (along with
   :exc:`asyncio.QueueShutDown`) for queue termination.
   (Contributed by Laurie Opperman and Yves Duprat in :gh:`104228`.)
diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py
index 57f0123..f2ee964 100644
--- a/Lib/asyncio/taskgroups.py
+++ b/Lib/asyncio/taskgroups.py
@@ -77,12 +77,6 @@ class TaskGroup:
             propagate_cancellation_error = exc
         else:
             propagate_cancellation_error = None
-        if self._parent_cancel_requested:
-            # If this flag is set we *must* call uncancel().
-            if self._parent_task.uncancel() == 0:
-                # If there are no pending cancellations left,
-                # don't propagate CancelledError.
-                propagate_cancellation_error = None
 
         if et is not None:
             if not self._aborting:
@@ -130,6 +124,13 @@ class TaskGroup:
         if self._base_error is not None:
             raise self._base_error
 
+        if self._parent_cancel_requested:
+            # If this flag is set we *must* call uncancel().
+            if self._parent_task.uncancel() == 0:
+                # If there are no pending cancellations left,
+                # don't propagate CancelledError.
+                propagate_cancellation_error = None
+
         # Propagate CancelledError if there is one, except if there
         # are other errors -- those have priority.
         if propagate_cancellation_error is not None and not self._errors:
@@ -139,6 +140,12 @@ class TaskGroup:
             self._errors.append(exc)
 
         if self._errors:
+            # If the parent task is being cancelled from the outside
+            # of the taskgroup, un-cancel and re-cancel the parent task,
+            # which will keep the cancel count stable.
+            if self._parent_task.cancelling():
+                self._parent_task.uncancel()
+                self._parent_task.cancel()
             # Exceptions are heavy objects that can have object
             # cycles (bad for GC); let's not keep a reference to
             # a bunch of them.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 7fb697b9..dadcb5b 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -255,6 +255,8 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
         """
         if self._num_cancels_requested > 0:
             self._num_cancels_requested -= 1
+            if self._num_cancels_requested == 0:
+                self._must_cancel = False
         return self._num_cancels_requested
 
     def __eager_start(self):
diff --git a/Lib/test/test_asyncio/test_taskgroups.py b/Lib/test/test_asyncio/test_taskgroups.py
index 1ec8116..4852536 100644
--- a/Lib/test/test_asyncio/test_taskgroups.py
+++ b/Lib/test/test_asyncio/test_taskgroups.py
@@ -833,6 +833,72 @@ class TestTaskGroup(unittest.IsolatedAsyncioTestCase):
         loop = asyncio.get_event_loop()
         loop.run_until_complete(run_coro_after_tg_closes())
 
+    async def test_cancelling_level_preserved(self):
+        async def raise_after(t, e):
+            await asyncio.sleep(t)
+            raise e()
+
+        try:
+            async with asyncio.TaskGroup() as tg:
+                tg.create_task(raise_after(0.0, RuntimeError))
+        except* RuntimeError:
+            pass
+        self.assertEqual(asyncio.current_task().cancelling(), 0)
+
+    async def test_nested_groups_both_cancelled(self):
+        async def raise_after(t, e):
+            await asyncio.sleep(t)
+            raise e()
+
+        try:
+            async with asyncio.TaskGroup() as outer_tg:
+                try:
+                    async with asyncio.TaskGroup() as inner_tg:
+                        inner_tg.create_task(raise_after(0, RuntimeError))
+                        outer_tg.create_task(raise_after(0, ValueError))
+                except* RuntimeError:
+                    pass
+                else:
+                    self.fail("RuntimeError not raised")
+            self.assertEqual(asyncio.current_task().cancelling(), 1)
+        except* ValueError:
+            pass
+        else:
+            self.fail("ValueError not raised")
+        self.assertEqual(asyncio.current_task().cancelling(), 0)
+
+    async def test_error_and_cancel(self):
+        event = asyncio.Event()
+
+        async def raise_error():
+            event.set()
+            await asyncio.sleep(0)
+            raise RuntimeError()
+
+        async def inner():
+            try:
+                async with taskgroups.TaskGroup() as tg:
+                    tg.create_task(raise_error())
+                    await asyncio.sleep(1)
+                    self.fail("Sleep in group should have been cancelled")
+            except* RuntimeError:
+                self.assertEqual(asyncio.current_task().cancelling(), 1)
+            self.assertEqual(asyncio.current_task().cancelling(), 1)
+            await asyncio.sleep(1)
+            self.fail("Sleep after group should have been cancelled")
+
+        async def outer():
+            t = asyncio.create_task(inner())
+            await event.wait()
+            self.assertEqual(t.cancelling(), 0)
+            t.cancel()
+            self.assertEqual(t.cancelling(), 1)
+            with self.assertRaises(asyncio.CancelledError):
+                await t
+            self.assertTrue(t.cancelled())
+
+        await outer()
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index bc6d88e..5b09c81 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -684,6 +684,30 @@ class BaseTaskTests:
         finally:
             loop.close()
 
+    def test_uncancel_resets_must_cancel(self):
+
+        async def coro():
+            await fut
+            return 42
+
+        loop = asyncio.new_event_loop()
+        fut = asyncio.Future(loop=loop)
+        task = self.new_task(loop, coro())
+        loop.run_until_complete(asyncio.sleep(0))  # Get task waiting for fut
+        fut.set_result(None)  # Make task runnable
+        try:
+            task.cancel()  # Enter cancelled state
+            self.assertEqual(task.cancelling(), 1)
+            self.assertTrue(task._must_cancel)
+
+            task.uncancel()  # Undo cancellation
+            self.assertEqual(task.cancelling(), 0)
+            self.assertFalse(task._must_cancel)
+        finally:
+            res = loop.run_until_complete(task)
+            self.assertEqual(res, 42)
+            loop.close()
+
     def test_cancel(self):
 
         def gen():
diff --git a/Misc/NEWS.d/next/Library/2024-04-04-15-28-12.gh-issue-116720.aGhXns.rst b/Misc/NEWS.d/next/Library/2024-04-04-15-28-12.gh-issue-116720.aGhXns.rst
new file mode 100644
index 0000000..39c7d6b
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-04-04-15-28-12.gh-issue-116720.aGhXns.rst
@@ -0,0 +1,18 @@
+Improved behavior of :class:`asyncio.TaskGroup` when an external cancellation
+collides with an internal cancellation. For example, when two task groups
+are nested and both experience an exception in a child task simultaneously,
+it was possible that the outer task group would misbehave, because
+its internal cancellation was swallowed by the inner task group.
+
+In the case where a task group is cancelled externally and also must
+raise an :exc:`ExceptionGroup`, it will now call the parent task's
+:meth:`~asyncio.Task.cancel` method. This ensures that a
+:exc:`asyncio.CancelledError` will be raised at the next
+:keyword:`await`, so the cancellation is not lost.
+
+An added benefit of these changes is that task groups now preserve the
+cancellation count (:meth:`asyncio.Task.cancelling`).
+
+In order to handle some corner cases, :meth:`asyncio.Task.uncancel` may now
+reset the undocumented ``_must_cancel`` flag when the cancellation count
+reaches zero.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 29246cf..b886051 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -2393,6 +2393,9 @@ _asyncio_Task_uncancel_impl(TaskObj *self)
 {
     if (self->task_num_cancels_requested > 0) {
         self->task_num_cancels_requested -= 1;
+        if (self->task_num_cancels_requested == 0) {
+            self->task_must_cancel = 0;
+        }
     }
     return PyLong_FromLong(self->task_num_cancels_requested);
 }