[Async-sig] pattern for handling interrupt signals in asyncio

[Chris Jerdonek]

I want to share a pattern I came up with for handling interrupt
signals in asyncio to see if you had any feedback (ways to make it
easier, similar approaches, etc).

I wanted something that was easy to check and reason about. I'm
already familiar with some of the pitfalls in handling signals, for
example as described in Nathaniel's Control-C blog post announced
here:
https://mail.python.org/pipermail/async-sig/2017-April/thread.html

The basic idea is to create a Future to run alongside the main
coroutine whose only purpose is to "catch" the signal. And then call--

    asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED)

When a signal is received, both tasks stop, and then you have access
to the main task (which will be pending) for things like cleanup and
inspection.

One advantage of this approach is that it lets you put all your
cleanup logic in the main program instead of putting some of it in the
signal handler. You also don't need to worry about things like
handling KeyboardInterrupt at arbitrary points in your code.

I'm including the code at bottom.

On the topic of asyncio.run() that I mentioned in an earlier email
[1], it doesn't look like the run() API posted in PR #465 [2] has
hooks to support what I'm describing (but I could be wrong). So maybe
this is another use case that the future API should contemplate.

--Chris

[1] https://mail.python.org/pipermail/async-sig/2017-August/000373.html
[2] https://github.com/python/asyncio/pull/465


import asyncio
import io
import signal

def _cleanup(loop):
    try:
        loop.run_until_complete(loop.shutdown_asyncgens())
    finally:
        loop.close()

def handle_sigint(future):
    future.set_result(signal.SIGINT)

async def run():
    print('running...')
    await asyncio.sleep(1000000)

def get_message(sig, task):
    stream = io.StringIO()
    task.print_stack(file=stream)
    traceback = stream.getvalue()
    return f'interrupted by {sig.name}:\n{traceback}'

def main(coro):
    loop = asyncio.new_event_loop()

    try:
        # This is made truthy if the loop is interrupted by a signal.
        interrupted = []

        future = asyncio.Future(loop=loop)
        future.add_done_callback(lambda future: interrupted.append(1))

        loop.add_signal_handler(signal.SIGINT, handle_sigint, future)

        futures = [future, coro]
        future = asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED)
        done, pending = loop.run_until_complete(future)

        if interrupted:
            # Do whatever cleanup you want here and/or get the stacktrace
            # of the interrupted main task.
            sig = done.pop().result()
            task = pending.pop()
            msg = get_message(sig, task)

            task.cancel()
            raise KeyboardInterrupt(msg)

    finally:
        _cleanup(loop)

main(run())


Below is what the code above outputs if you run it and then press Control-C:

running...
^CTraceback (most recent call last):
  File "test-signal.py", line 54, in <module>
    main(run())
  File "test-signal.py", line 49, in main
    raise KeyboardInterrupt(msg)
KeyboardInterrupt: interrupted by SIGINT:
Stack for <Task pending coro=<run() running at test-signal.py:17>
wait_for=<Future pending cb=[<TaskWakeupMethWrapper object at
0x10fe9b9a8>()]>> (most recent call last):
  File "test-signal.py", line 17, in run
    await asyncio.sleep(1000000)