From 7c1f069d9e786b79d59737943effb9f2e5ef0832 Mon Sep 17 00:00:00 2001 From: Connor Maynes Date: Wed, 30 Nov 2022 08:01:34 -0500 Subject: [PATCH 1/5] 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 571013745aa03a..7c56949b5edfb7 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 Date: Wed, 30 Nov 2022 08:08:42 -0500 Subject: [PATCH 2/5] 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 7c56949b5edfb7..b74d0aa8c6f464 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. From db41d7f849e7b053bb79a54dc6b2498dcc454a9d Mon Sep 17 00:00:00 2001 From: Connor Maynes Date: Wed, 30 Nov 2022 09:00:30 -0500 Subject: [PATCH 3/5] update NEWS --- .../next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst | 3 +++ 1 file changed, 3 insertions(+) create mode 100644 Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst diff --git a/Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst b/Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst new file mode 100644 index 00000000000000..76f6cdea13c71f --- /dev/null +++ b/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. From d1710a7a1953167ba245d5c70b2c153cfc795d47 Mon Sep 17 00:00:00 2001 From: Connor Maynes Date: Wed, 25 Jan 2023 14:08:40 -0500 Subject: [PATCH 4/5] turn shields task property into int and remove shielded decorator --- Lib/asyncio/tasks.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index b74d0aa8c6f464..c1855083e0cef7 100644 --- a/Lib/asyncio/tasks.py +++ b/Lib/asyncio/tasks.py @@ -78,29 +78,18 @@ def _set_task_name(task, 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) + task.shields += 1 try: yield finally: - if task.shields.pop() is not shield: + task.shields -= 1 + if task.shields < 0: raise RuntimeError( f"Unbalanced enter/exit of shield scopes for task {task!r}" ) -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. @@ -121,7 +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 = [] + self.shields = 0 if self._source_traceback: del self._source_traceback[-1] if not coroutines.iscoroutine(coro): From 47943223ed88ac6146990f807d8aa4ee8f509ca6 Mon Sep 17 00:00:00 2001 From: Connor Maynes Date: Wed, 25 Jan 2023 14:20:23 -0500 Subject: [PATCH 5/5] fix: expose shield scope in __all__ --- Lib/asyncio/tasks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py index c1855083e0cef7..3c4b2e265869a4 100644 --- a/Lib/asyncio/tasks.py +++ b/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', )