Possible asyncio.Task cancellation fix.

QuantumTim · QuantumTim · commit b6da21bf5f0d · 2022-01-04T18:29:58.000Z
- Ensures CancelledError can't be raised at the expense of leaking other values/results.
- Comes at the expense of being able to (easily) abort cancellations.
  - You can still catch and ignore CancelledError, but it will keep re-appearing.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
@@ -201,25 +201,25 @@ def cancel(self, msg=None):
         self._log_traceback = False
         if self.done():
             return False
+        # Once set, we keep trying to cancel the Task at every
+        # safe opportunity. There's currently no way to stop this,
+        # but maybe we could add CancelledError.abort() or something
+        # for tasks to explicitly stop cancellation.
+        self._must_cancel = 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
                 # catches and ignores the cancellation so we may have
                 # to cancel it again later.
                 return True
         # It must be the case that self.__step is already scheduled.
-        self._must_cancel = True
         self._cancel_message = msg
         return True
 
     def __step(self, exc=None):
         if self.done():
             raise exceptions.InvalidStateError(
                 f'_step(): already done: {self!r}, {exc!r}')
-        if self._must_cancel:
-            if not isinstance(exc, exceptions.CancelledError):
-                exc = self._make_cancelled_error()
-            self._must_cancel = False
         coro = self._coro
         self._fut_waiter = None
 
@@ -234,11 +234,13 @@ def __step(self, exc=None):
                 result = coro.throw(exc)
         except StopIteration as exc:
             if self._must_cancel:
-                # Task is cancelled right before coro stops.
+                # We were cancelled, but we're done, so we just pretend
+                # that the cancellation came in too late and the result
+                # was already ready. If the calling Task was the source
+                # of the cancellation, then its own _must_cancel will
+                # still be set, so it will still keep trying to cancel.
                 self._must_cancel = False
-                super().cancel(msg=self._cancel_message)
-            else:
-                super().set_result(exc.value)
+            super().set_result(exc.value)
         except exceptions.CancelledError as exc:
             # Save the original exception so we can chain it later.
             self._cancelled_exc = exc
@@ -272,7 +274,12 @@ def __step(self, exc=None):
                         if self._must_cancel:
                             if self._fut_waiter.cancel(
                                     msg=self._cancel_message):
-                                self._must_cancel = False
+                                # Even if we cancelled the future, we
+                                # leave _must_cancel set because the
+                                # future might racily return a result
+                                # at the same time as the cancellation,
+                                # so we'll need to retry cancellation.
+                                pass
                 else:
                     new_exc = RuntimeError(
                         f'yield was used instead of yield from '
