Skip to content

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

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 2 commits into from
Closed

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

Changes from 1 commit
Commits
Show all changes
2 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
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
Next Next commit
gh-99714: add shield scope to asyncio
  • Loading branch information
Connor Maynes
Connor Maynes committed Nov 30, 2022
commit 7c1f069d9e786b79d59737943effb9f2e5ef0832
23 changes: 20 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,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.

Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down