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).

Author: Tinche

Lib/asyncio/tasks.py

@@ -105,7 +105,7 @@ def __init__(self, coro, *, loop=None, name=None):
         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 @@ def cancel(self, msg=None):
         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 @@ def cancel(self, msg=None):
         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():

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,22 +2191,24 @@ 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;
 
     if (self->task_state != STATE_PENDING) {
         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]