Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
whatamithinking wants to merge 5 commits into from
Closed

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7c1f069
gh-99714: add shield scope to asyncio
Nov 30, 2022
9de6823
gh-99714: add shielded decorator
Nov 30, 2022
db41d7f
update NEWS
Nov 30, 2022
d1710a7
turn shields task property into int and remove shielded decorator
Jan 25, 2023
4794322
fix: expose shield scope in __all__
Jan 25, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 30 additions & 3 deletions Lib/asyncio/tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
import types
import warnings
import weakref
import contextlib
from types import GenericAlias

from . import base_tasks
Expand Down Expand Up @@ -75,6 +76,31 @@ 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

@gvanrossum gvanrossum Jan 20, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But who would do that? Let them suffer. Task().shields can just be an integer counter.

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}"
)


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

@gvanrossum gvanrossum Jan 20, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we absolutely, truly need two ways to do it (a decorator and a context manager)? I'd start with just the context manager.



class Task(futures._PyFuture): # Inherit Python Task implementation
# from a Python Future implementation.

Expand All @@ -95,6 +121,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):
Expand Down Expand Up @@ -219,7 +246,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
Expand Down Expand Up @@ -254,7 +281,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
Expand Down Expand Up @@ -307,7 +334,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
Expand Down
3 changes: 3 additions & 0 deletions Misc/NEWS.d/next/Library/2022-11-30-09-00-13.gh-issue-99714.Xkbibc.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.