Add an API to Task to manage 'cancel_requested' flag

gvanrossum · gvanrossum · commit 77ec0e42b3fc · 2022-02-14T11:12:43.000-08:00
This means we no longer have to monkey-patch the parent task. It does introduce new semantics for task cancellation: When a task is cancelled, further attempts to cancel it have *no* effect unless the task calls self.uncancel(). Borrowed from pythonGH-31313 by @asvetlov.
diff --git a/Lib/asyncio/taskgroups.py b/Lib/asyncio/taskgroups.py
@@ -78,7 +78,6 @@ async def __aenter__(self):
         if self._parent_task is None:
             raise RuntimeError(
                 f'TaskGroup {self!r} cannot determine the parent task')
-        self._patch_task(self._parent_task)
 
         return self
 
@@ -95,7 +94,7 @@ async def __aexit__(self, et, exc, tb):
             if self._parent_cancel_requested:
                 # Only if we did request task to cancel ourselves
                 # we mark it as no longer cancelled.
-                self._parent_task.__cancel_requested__ = False
+                self._parent_task.uncancel()
             else:
                 propagate_cancellation_error = et
 
@@ -186,26 +185,6 @@ def _is_base_error(self, exc: BaseException) -> bool:
         assert isinstance(exc, BaseException)
         return isinstance(exc, (SystemExit, KeyboardInterrupt))
 
-    def _patch_task(self, task):
-        # In the future we'll need proper API on asyncio.Task to
-        # make TaskGroups possible. We need to be able to access
-        # information about task cancellation, more specifically,
-        # we need a flag to say if a task was cancelled or not.
-        # We also need to be able to flip that flag.
-
-        def _task_cancel(self, msg=None):
-            self.__cancel_requested__ = True
-            return asyncio.Task.cancel(self, msg)
-
-        if hasattr(task, '__cancel_requested__'):
-            return
-
-        task.__cancel_requested__ = False
-        # confirm that we were successful at adding the new attribute:
-        assert not task.__cancel_requested__
-
-        task.cancel = types.MethodType(_task_cancel, task)
-
     def _abort(self):
         self._aborting = True
 
@@ -244,7 +223,7 @@ def _on_task_done(self, task):
             return
 
         self._abort()
-        if not self._parent_task.__cancel_requested__:
+        if not self._parent_task.cancelling():
             # If parent task *is not* being cancelled, it means that we want
             # to manually cancel it to abort whatever is being run right now
             # in the TaskGroup.  But we want to mark parent task as
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -105,6 +105,7 @@ def __init__(self, coro, *, loop=None, name=None):
         else:
             self._name = str(name)
 
+        self._cancel_requested = False
         self._must_cancel = False
         self._fut_waiter = None
         self._coro = coro
@@ -201,6 +202,9 @@ def cancel(self, msg=None):
         self._log_traceback = False
         if self.done():
             return False
+        if self._cancel_requested:
+            return False
+        self._cancel_requested = True
         if self._fut_waiter is not None:
             if self._fut_waiter.cancel(msg=msg):
                 # Leave self._fut_waiter; it may be a Task that
@@ -212,6 +216,16 @@ def cancel(self, msg=None):
         self._cancel_message = msg
         return True
 
+    def cancelling(self):
+        return self._cancel_requested
+
+    def uncancel(self):
+        if self._cancel_requested:
+            self._cancel_requested = False
+            return True
+        else:
+            return False
+
     def __step(self, exc=None):
         if self.done():
             raise exceptions.InvalidStateError(
@@ -634,7 +648,7 @@ def _ensure_future(coro_or_future, *, loop=None):
         loop = events._get_event_loop(stacklevel=4)
     try:
         return loop.create_task(coro_or_future)
-    except RuntimeError: 
+    except RuntimeError:
         if not called_wrap_awaitable:
             coro_or_future.close()
         raise
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -90,6 +90,7 @@ typedef struct {
     PyObject *task_context;
     int task_must_cancel;
     int task_log_destroy_pending;
+    int task_cancel_requested;
 } TaskObj;
 
 typedef struct {
@@ -2038,6 +2039,7 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
     Py_CLEAR(self->task_fut_waiter);
     self->task_must_cancel = 0;
     self->task_log_destroy_pending = 1;
+    self->task_cancel_requested = 0;
     Py_INCREF(coro);
     Py_XSETREF(self->task_coro, coro);
 
@@ -2204,6 +2206,11 @@ _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
         Py_RETURN_FALSE;
     }
 
+    if (self->task_cancel_requested) {
+        Py_RETURN_FALSE;
+    }
+    self->task_cancel_requested = 1;
+
     if (self->task_fut_waiter) {
         PyObject *res;
         int is_true;
@@ -2231,6 +2238,56 @@ _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
     Py_RETURN_TRUE;
 }
 
+/*[clinic input]
+_asyncio.Task.cancelling
+
+Return True if the task is in the process of being cancelled.
+
+This is set once .cancel() is called
+and remains set until .uncancel() is called.
+
+As long as this flag is set, further .cancel() calls will be ignored,
+until .uncancel() is called to reset it.
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task_cancelling_impl(TaskObj *self)
+/*[clinic end generated code: output=803b3af96f917d7e input=c50e50f9c3ca4676]*/
+/*[clinic end generated code]*/
+{
+    if (self->task_cancel_requested) {
+        Py_RETURN_TRUE;
+    }
+    else {
+        Py_RETURN_FALSE;
+    }
+}
+
+/*[clinic input]
+_asyncio.Task.uncancel
+
+Reset the flag returned by cancelling().
+
+This should be used by tasks that catch CancelledError
+and wish to continue indefinitely until they are cancelled again.
+
+Returns the previous value of the flag.
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task_uncancel_impl(TaskObj *self)
+/*[clinic end generated code: output=58184d236a817d3c input=5db95e28fcb6f7cd]*/
+/*[clinic end generated code]*/
+{
+    if (self->task_cancel_requested) {
+        self->task_cancel_requested = 0;
+        Py_RETURN_TRUE;
+    }
+    else {
+        Py_RETURN_FALSE;
+    }
+}
+
 /*[clinic input]
 _asyncio.Task.get_stack
 
@@ -2454,6 +2511,8 @@ static PyMethodDef TaskType_methods[] = {
     _ASYNCIO_TASK_SET_RESULT_METHODDEF
     _ASYNCIO_TASK_SET_EXCEPTION_METHODDEF
     _ASYNCIO_TASK_CANCEL_METHODDEF
+    _ASYNCIO_TASK_CANCELLING_METHODDEF
+    _ASYNCIO_TASK_UNCANCEL_METHODDEF
     _ASYNCIO_TASK_GET_STACK_METHODDEF
     _ASYNCIO_TASK_PRINT_STACK_METHODDEF
     _ASYNCIO_TASK__MAKE_CANCELLED_ERROR_METHODDEF
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
