bpo-46771: Implement task cancel requests counter (GH-31513)

This changes cancelling() and uncancel() to return the count of pending cancellations.

This can be used to avoid bugs in certain edge cases (e.g. two timeouts going off at the same time).

3 files changed, 48 insertions, 46 deletions

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 25a650f..38c2385 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -105,7 +105,7 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
         else:
             self._name = str(name)
 
-        self._cancel_requested = False
+        self._num_cancels_requested = 0
         self._must_cancel = False
         self._fut_waiter = None
         self._coro = coro
@@ -198,13 +198,15 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
         task will be marked as cancelled when the wrapped coroutine
         terminates with a CancelledError exception (even if cancel()
         was not called).
+
+        This also increases the task's count of cancellation requests.
         """
         self._log_traceback = False
         if self.done():
             return False
-        if self._cancel_requested:
+        self._num_cancels_requested += 1
+        if self._num_cancels_requested > 1:
             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
@@ -217,14 +219,24 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
         return True
 
     def cancelling(self):
-        return self._cancel_requested
+        """Return the count of the task's cancellation requests.
+
+        This count is incremented when .cancel() is called
+        and may be decremented using .uncancel().
+        """
+        return self._num_cancels_requested
 
     def uncancel(self):
-        if self._cancel_requested:
-            self._cancel_requested = False
-            return True
-        else:
-            return False
+        """Decrement the task's count of cancellation requests.
+
+        This should be used by tasks that catch CancelledError
+        and wish to continue indefinitely until they are cancelled again.
+
+        Returns the remaining number of cancellation requests.
+        """
+        if self._num_cancels_requested > 0:
+            self._num_cancels_requested -= 1
+        return self._num_cancels_requested
 
     def __step(self, exc=None):
         if self.done():
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index 0d6b21d..8847123 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -91,7 +91,7 @@ typedef struct {
     PyObject *task_context;
     int task_must_cancel;
     int task_log_destroy_pending;
-    int task_cancel_requested;
+    int task_num_cancels_requested;
 } TaskObj;
 
 typedef struct {
@@ -2036,7 +2036,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;
+    self->task_num_cancels_requested = 0;
     Py_INCREF(coro);
     Py_XSETREF(self->task_coro, coro);
 
@@ -2191,11 +2191,13 @@ not return True (unless the task was already cancelled).  A
 task will be marked as cancelled when the wrapped coroutine
 terminates with a CancelledError exception (even if cancel()
 was not called).
+
+This also increases the task's count of cancellation requests.
 [clinic start generated code]*/
 
 static PyObject *
 _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
-/*[clinic end generated code: output=c66b60d41c74f9f1 input=f4ff8e8ffc5f1c00]*/
+/*[clinic end generated code: output=c66b60d41c74f9f1 input=7bb51bf25974c783]*/
 {
     self->task_log_tb = 0;
 
@@ -2203,10 +2205,10 @@ _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
         Py_RETURN_FALSE;
     }
 
-    if (self->task_cancel_requested) {
+    self->task_num_cancels_requested += 1;
+    if (self->task_num_cancels_requested > 1) {
         Py_RETURN_FALSE;
     }
-    self->task_cancel_requested = 1;
 
     if (self->task_fut_waiter) {
         PyObject *res;
@@ -2238,51 +2240,40 @@ _asyncio_Task_cancel_impl(TaskObj *self, PyObject *msg)
 /*[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.
+Return the count of the task's cancellation requests.
 
-As long as this flag is set, further .cancel() calls will be ignored,
-until .uncancel() is called to reset it.
+This count is incremented when .cancel() is called
+and may be decremented using .uncancel().
 [clinic start generated code]*/
 
 static PyObject *
 _asyncio_Task_cancelling_impl(TaskObj *self)
-/*[clinic end generated code: output=803b3af96f917d7e input=c50e50f9c3ca4676]*/
+/*[clinic end generated code: output=803b3af96f917d7e input=b625224d310cbb17]*/
 /*[clinic end generated code]*/
 {
-    if (self->task_cancel_requested) {
-        Py_RETURN_TRUE;
-    }
-    else {
-        Py_RETURN_FALSE;
-    }
+    return PyLong_FromLong(self->task_num_cancels_requested);
 }
 
 /*[clinic input]
 _asyncio.Task.uncancel
 
-Reset the flag returned by cancelling().
+Decrement the task's count of cancellation requests.
 
 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.
+Returns the remaining number of cancellation requests.
 [clinic start generated code]*/
 
 static PyObject *
 _asyncio_Task_uncancel_impl(TaskObj *self)
-/*[clinic end generated code: output=58184d236a817d3c input=5db95e28fcb6f7cd]*/
+/*[clinic end generated code: output=58184d236a817d3c input=68f81a4b90b46be2]*/
 /*[clinic end generated code]*/
 {
-    if (self->task_cancel_requested) {
-        self->task_cancel_requested = 0;
-        Py_RETURN_TRUE;
-    }
-    else {
-        Py_RETURN_FALSE;
+    if (self->task_num_cancels_requested > 0) {
+        self->task_num_cancels_requested -= 1;
     }
+    return PyLong_FromLong(self->task_num_cancels_requested);
 }
 
 /*[clinic input]
diff --git a/Modules/clinic/_asynciomodule.c.h b/Modules/clinic/_asynciomodule.c.h
index 5648e14..2b84ef0 100644
--- a/Modules/clinic/_asynciomodule.c.h
+++ b/Modules/clinic/_asynciomodule.c.h
@@ -414,7 +414,9 @@ PyDoc_STRVAR(_asyncio_Task_cancel__doc__,
 "not return True (unless the task was already cancelled).  A\n"
 "task will be marked as cancelled when the wrapped coroutine\n"
 "terminates with a CancelledError exception (even if cancel()\n"
-"was not called).");
+"was not called).\n"
+"\n"
+"This also increases the task\'s count of cancellation requests.");
 
 #define _ASYNCIO_TASK_CANCEL_METHODDEF    \
     {"cancel", (PyCFunction)(void(*)(void))_asyncio_Task_cancel, METH_FASTCALL|METH_KEYWORDS, _asyncio_Task_cancel__doc__},
@@ -451,13 +453,10 @@ PyDoc_STRVAR(_asyncio_Task_cancelling__doc__,
 "cancelling($self, /)\n"
 "--\n"
 "\n"
-"Return True if the task is in the process of being cancelled.\n"
-"\n"
-"This is set once .cancel() is called\n"
-"and remains set until .uncancel() is called.\n"
+"Return the count of the task\'s cancellation requests.\n"
 "\n"
-"As long as this flag is set, further .cancel() calls will be ignored,\n"
-"until .uncancel() is called to reset it.");
+"This count is incremented when .cancel() is called\n"
+"and may be decremented using .uncancel().");
 
 #define _ASYNCIO_TASK_CANCELLING_METHODDEF    \
     {"cancelling", (PyCFunction)_asyncio_Task_cancelling, METH_NOARGS, _asyncio_Task_cancelling__doc__},
@@ -475,12 +474,12 @@ PyDoc_STRVAR(_asyncio_Task_uncancel__doc__,
 "uncancel($self, /)\n"
 "--\n"
 "\n"
-"Reset the flag returned by cancelling().\n"
+"Decrement the task\'s count of cancellation requests.\n"
 "\n"
 "This should be used by tasks that catch CancelledError\n"
 "and wish to continue indefinitely until they are cancelled again.\n"
 "\n"
-"Returns the previous value of the flag.");
+"Returns the remaining number of cancellation requests.");
 
 #define _ASYNCIO_TASK_UNCANCEL_METHODDEF    \
     {"uncancel", (PyCFunction)_asyncio_Task_uncancel, METH_NOARGS, _asyncio_Task_uncancel__doc__},
@@ -918,4 +917,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=c02708a9d6a774cc input=a9049054013a1b77]*/
+/*[clinic end generated code: output=344927e9b6016ad7 input=a9049054013a1b77]*/