bpo-32314: Fix asyncio.run() to cancel runinng tasks on shutdown (#5262)

This commit is contained in:
Yury Selivanov 2018-01-21 14:56:59 -05:00 committed by GitHub
parent fc2f407829
commit a4afcdfa55
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 122 additions and 15 deletions

View file

@ -228,14 +228,9 @@ class BaseEventLoop(events.AbstractEventLoop):
self._coroutine_origin_tracking_enabled = False
self._coroutine_origin_tracking_saved_depth = None
if hasattr(sys, 'get_asyncgen_hooks'):
# Python >= 3.6
# A weak set of all asynchronous generators that are
# being iterated by the loop.
self._asyncgens = weakref.WeakSet()
else:
self._asyncgens = None
# A weak set of all asynchronous generators that are
# being iterated by the loop.
self._asyncgens = weakref.WeakSet()
# Set to True when `loop.shutdown_asyncgens` is called.
self._asyncgens_shutdown_called = False
@ -354,7 +349,7 @@ class BaseEventLoop(events.AbstractEventLoop):
"""Shutdown all active asynchronous generators."""
self._asyncgens_shutdown_called = True
if self._asyncgens is None or not len(self._asyncgens):
if not len(self._asyncgens):
# If Python version is <3.6 or we don't have any asynchronous
# generators alive.
return
@ -386,10 +381,10 @@ class BaseEventLoop(events.AbstractEventLoop):
'Cannot run the event loop while another loop is running')
self._set_coroutine_origin_tracking(self._debug)
self._thread_id = threading.get_ident()
if self._asyncgens is not None:
old_agen_hooks = sys.get_asyncgen_hooks()
sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
finalizer=self._asyncgen_finalizer_hook)
old_agen_hooks = sys.get_asyncgen_hooks()
sys.set_asyncgen_hooks(firstiter=self._asyncgen_firstiter_hook,
finalizer=self._asyncgen_finalizer_hook)
try:
events._set_running_loop(self)
while True:
@ -401,8 +396,7 @@ class BaseEventLoop(events.AbstractEventLoop):
self._thread_id = None
events._set_running_loop(None)
self._set_coroutine_origin_tracking(False)
if self._asyncgens is not None:
sys.set_asyncgen_hooks(*old_agen_hooks)
sys.set_asyncgen_hooks(*old_agen_hooks)
def run_until_complete(self, future):
"""Run until the Future is done.
@ -1374,6 +1368,7 @@ class BaseEventLoop(events.AbstractEventLoop):
- 'message': Error message;
- 'exception' (optional): Exception object;
- 'future' (optional): Future instance;
- 'task' (optional): Task instance;
- 'handle' (optional): Handle instance;
- 'protocol' (optional): Protocol instance;
- 'transport' (optional): Transport instance;

View file

@ -2,6 +2,7 @@ __all__ = 'run',
from . import coroutines
from . import events
from . import tasks
def run(main, *, debug=False):
@ -42,7 +43,31 @@ def run(main, *, debug=False):
return loop.run_until_complete(main)
finally:
try:
_cancel_all_tasks(loop)
loop.run_until_complete(loop.shutdown_asyncgens())
finally:
events.set_event_loop(None)
loop.close()
def _cancel_all_tasks(loop):
to_cancel = [task for task in tasks.all_tasks(loop)
if not task.done()]
if not to_cancel:
return
for task in to_cancel:
task.cancel()
loop.run_until_complete(
tasks.gather(*to_cancel, loop=loop, return_exceptions=True))
for task in to_cancel:
if task.cancelled():
continue
if task.exception() is not None:
loop.call_exception_handler({
'message': 'unhandled exception during asyncio.run() shutdown',
'exception': task.exception(),
'task': task,
})