gh-99714: asyncio - add shield_scope context manager and shielded decorator

Lib/asyncio/tasks.py

@@ -4,7 +4,7 @@
     'Task', 'create_task',
     'FIRST_COMPLETED', 'FIRST_EXCEPTION', 'ALL_COMPLETED',
     'wait', 'wait_for', 'as_completed', 'sleep',
-    'gather', 'shield', 'ensure_future', 'run_coroutine_threadsafe',
+    'gather', 'shield', 'shield_scope', 'ensure_future', 'run_coroutine_threadsafe',
     'current_task', 'all_tasks',
     '_register_task', '_unregister_task', '_enter_task', '_leave_task',
 )
@@ -17,6 +17,7 @@
 import types
 import warnings
 import weakref
+import contextlib
 from types import GenericAlias
 
 from . import base_tasks
@@ -75,6 +76,20 @@ def _set_task_name(task, name):
             set_name(name)
 
 
+@contextlib.contextmanager
+def shield_scope():
+    task = current_task()
+    task.shields += 1
+    try:
+        yield
+    finally:
+        task.shields -= 1
+        if task.shields < 0:
+            raise RuntimeError(
+                f"Unbalanced enter/exit of shield scopes for task {task!r}"
+            )
+
+
 class Task(futures._PyFuture):  # Inherit Python Task implementation
                                 # from a Python Future implementation.
 
@@ -95,6 +110,7 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
 
     def __init__(self, coro, *, loop=None, name=None, context=None):
         super().__init__(loop=loop)
+        self.shields = 0
         if self._source_traceback:
             del self._source_traceback[-1]
         if not coroutines.iscoroutine(coro):
@@ -219,7 +235,7 @@ def cancel(self, msg=None):
         # Also remember that this is duplicated in _asynciomodule.c.
         # if self._num_cancels_requested > 1:
         #     return False
-        if self._fut_waiter is not None:
+        if self._fut_waiter is not None and not self.shields:
             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
@@ -254,7 +270,7 @@ def __step(self, exc=None):
         if self.done():
             raise exceptions.InvalidStateError(
                 f'_step(): already done: {self!r}, {exc!r}')
-        if self._must_cancel:
+        if self._must_cancel and not self.shields:
             if not isinstance(exc, exceptions.CancelledError):
                 exc = self._make_cancelled_error()
             self._must_cancel = False
@@ -307,7 +323,7 @@ def __step(self, exc=None):
                         result.add_done_callback(
                             self.__wakeup, context=self._context)
                         self._fut_waiter = result
-                        if self._must_cancel:
+                        if self._must_cancel and not self.shields:
                             if self._fut_waiter.cancel(
                                     msg=self._cancel_message):
                                 self._must_cancel = False

Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst

@@ -0,0 +1,3 @@
+Add :func:`shield_scope` context manager in :mod:`asyncio.tasks` to shield
+async blocks of code from cancellation. Add :func:`shielded` decorator in
+:mod:`asyncio.tasks` to shield functions from cancellation.