bpo-39622: Interrupt the main asyncio task on Ctrl+C (GH-32105)

Co-authored-by: Kumar Aditya <59607654+kumaraditya303@users.noreply.github.com>
This commit is contained in:
Andrew Svetlov 2022-03-30 15:15:06 +03:00 committed by GitHub
parent 04acfa94bb
commit f08a191882
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 122 additions and 2 deletions

View file

@ -119,3 +119,30 @@ Runner context manager
Embedded *loop* and *context* are created at the :keyword:`with` body entering Embedded *loop* and *context* are created at the :keyword:`with` body entering
or the first call of :meth:`run` or :meth:`get_loop`. or the first call of :meth:`run` or :meth:`get_loop`.
Handling Keyboard Interruption
==============================
.. versionadded:: 3.11
When :const:`signal.SIGINT` is raised by :kbd:`Ctrl-C`, :exc:`KeyboardInterrupt`
exception is raised in the main thread by default. However this doesn't work with
:mod:`asyncio` because it can interrupt asyncio internals and can hang the program from
exiting.
To mitigate this issue, :mod:`asyncio` handles :const:`signal.SIGINT` as follows:
1. :meth:`asyncio.Runner.run` installs a custom :const:`signal.SIGINT` handler before
any user code is executed and removes it when exiting from the function.
2. The :class:`~asyncio.Runner` creates the main task for the passed coroutine for its
execution.
3. When :const:`signal.SIGINT` is raised by :kbd:`Ctrl-C`, the custom signal handler
cancels the main task by calling :meth:`asyncio.Task.cancel` which raises
:exc:`asyncio.CancelledError` inside the the main task. This causes the Python stack
to unwind, ``try/except`` and ``try/finally`` blocks can be used for resource
cleanup. After the main task is cancelled, :meth:`asyncio.Runner.run` raises
:exc:`KeyboardInterrupt`.
4. A user could write a tight loop which cannot be interrupted by
:meth:`asyncio.Task.cancel`, in which case the second following :kbd:`Ctrl-C`
immediately raises the :exc:`KeyboardInterrupt` without cancelling the main task.

View file

@ -2,8 +2,13 @@ __all__ = ('Runner', 'run')
import contextvars import contextvars
import enum import enum
import functools
import threading
import signal
import sys
from . import coroutines from . import coroutines
from . import events from . import events
from . import exceptions
from . import tasks from . import tasks
@ -47,6 +52,7 @@ class Runner:
self._loop_factory = loop_factory self._loop_factory = loop_factory
self._loop = None self._loop = None
self._context = None self._context = None
self._interrupt_count = 0
def __enter__(self): def __enter__(self):
self._lazy_init() self._lazy_init()
@ -89,7 +95,28 @@ class Runner:
if context is None: if context is None:
context = self._context context = self._context
task = self._loop.create_task(coro, context=context) task = self._loop.create_task(coro, context=context)
if (threading.current_thread() is threading.main_thread()
and signal.getsignal(signal.SIGINT) is signal.default_int_handler
):
sigint_handler = functools.partial(self._on_sigint, main_task=task)
signal.signal(signal.SIGINT, sigint_handler)
else:
sigint_handler = None
self._interrupt_count = 0
try:
return self._loop.run_until_complete(task) return self._loop.run_until_complete(task)
except exceptions.CancelledError:
if self._interrupt_count > 0 and task.uncancel() == 0:
raise KeyboardInterrupt()
else:
raise # CancelledError
finally:
if (sigint_handler is not None
and signal.getsignal(signal.SIGINT) is sigint_handler
):
signal.signal(signal.SIGINT, signal.default_int_handler)
def _lazy_init(self): def _lazy_init(self):
if self._state is _State.CLOSED: if self._state is _State.CLOSED:
@ -105,6 +132,14 @@ class Runner:
self._context = contextvars.copy_context() self._context = contextvars.copy_context()
self._state = _State.INITIALIZED self._state = _State.INITIALIZED
def _on_sigint(self, signum, frame, main_task):
self._interrupt_count += 1
if self._interrupt_count == 1 and not main_task.done():
main_task.cancel()
# wakeup loop if it is blocked by select() with long timeout
self._loop.call_soon_threadsafe(lambda: None)
return
raise KeyboardInterrupt()
def run(main, *, debug=None): def run(main, *, debug=None):

View file

@ -1,7 +1,9 @@
import _thread
import asyncio import asyncio
import contextvars import contextvars
import gc import gc
import re import re
import threading
import unittest import unittest
from unittest import mock from unittest import mock
@ -12,6 +14,10 @@ def tearDownModule():
asyncio.set_event_loop_policy(None) asyncio.set_event_loop_policy(None)
def interrupt_self():
_thread.interrupt_main()
class TestPolicy(asyncio.AbstractEventLoopPolicy): class TestPolicy(asyncio.AbstractEventLoopPolicy):
def __init__(self, loop_factory): def __init__(self, loop_factory):
@ -298,7 +304,7 @@ class RunnerTests(BaseTest):
self.assertEqual(2, runner.run(get_context()).get(cvar)) self.assertEqual(2, runner.run(get_context()).get(cvar))
def test_recursine_run(self): def test_recursive_run(self):
async def g(): async def g():
pass pass
@ -318,6 +324,57 @@ class RunnerTests(BaseTest):
): ):
runner.run(f()) runner.run(f())
def test_interrupt_call_soon(self):
# The only case when task is not suspended by waiting a future
# or another task
assert threading.current_thread() is threading.main_thread()
async def coro():
with self.assertRaises(asyncio.CancelledError):
while True:
await asyncio.sleep(0)
raise asyncio.CancelledError()
with asyncio.Runner() as runner:
runner.get_loop().call_later(0.1, interrupt_self)
with self.assertRaises(KeyboardInterrupt):
runner.run(coro())
def test_interrupt_wait(self):
# interrupting when waiting a future cancels both future and main task
assert threading.current_thread() is threading.main_thread()
async def coro(fut):
with self.assertRaises(asyncio.CancelledError):
await fut
raise asyncio.CancelledError()
with asyncio.Runner() as runner:
fut = runner.get_loop().create_future()
runner.get_loop().call_later(0.1, interrupt_self)
with self.assertRaises(KeyboardInterrupt):
runner.run(coro(fut))
self.assertTrue(fut.cancelled())
def test_interrupt_cancelled_task(self):
# interrupting cancelled main task doesn't raise KeyboardInterrupt
assert threading.current_thread() is threading.main_thread()
async def subtask(task):
await asyncio.sleep(0)
task.cancel()
interrupt_self()
async def coro():
asyncio.create_task(subtask(asyncio.current_task()))
await asyncio.sleep(10)
with asyncio.Runner() as runner:
with self.assertRaises(asyncio.CancelledError):
runner.run(coro())
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()

View file

@ -0,0 +1 @@
Handle Ctrl+C in asyncio programs to interrupt the main task.