From 7c1f069d9e786b79d59737943effb9f2e5ef0832 Mon Sep 17 00:00:00 2001
From: Connor Maynes <connormaynes@gmail.com>
Date: Wed, 30 Nov 2022 08:01:34 -0500
Subject: [PATCH 1/2] gh-99714: add shield scope to asyncio

---
 Lib/asyncio/tasks.py | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 571013745aa03aa..7c56949b5edfb7d 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -17,6 +17,7 @@
 import types
 import warnings
 import weakref
+import contextlib
 from types import GenericAlias
 
 from . import base_tasks
@@ -75,6 +76,21 @@ def _set_task_name(task, name):
             set_name(name)
 
 
+@contextlib.contextmanager
+def shield_scope():
+    shield = object()  # just for making sure shields closed in order
+    # if __enter__ and __exit__ called directly
+    task = current_task()
+    task.shields.append(shield)
+    try:
+        yield
+    finally:
+        if task.shields.pop() is not shield:
+            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 +111,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 = []
         if self._source_traceback:
             del self._source_traceback[-1]
         if not coroutines.iscoroutine(coro):
@@ -219,7 +236,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 +271,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 +324,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

From 9de6823fe20ab6b89f3c7bb987eafd32af66dfff Mon Sep 17 00:00:00 2001
From: Connor Maynes <connormaynes@gmail.com>
Date: Wed, 30 Nov 2022 08:08:42 -0500
Subject: [PATCH 2/2] gh-99714: add shielded decorator

---
 Lib/asyncio/tasks.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 7c56949b5edfb7d..b74d0aa8c6f4643 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -91,6 +91,16 @@ def shield_scope():
             )
 
 
+def shielded(func):
+    """Decorator to shield the call of a async function or method,
+    wrapping it in a shield_scope()."""
+    @functools.wraps(func)
+    async def _shielded(*args, **kwargs):
+        with shield_scope():
+            return await func(*args, **kwargs)
+    return _shielded
+
+
 class Task(futures._PyFuture):  # Inherit Python Task implementation
                                 # from a Python Future implementation.