Asyncio: Choose early, choose wisely

Published: 2018-01-01

Tagged: python architecture programming guide

Over the past few months, I've had the pleasure to work on some cool asyncio-based code. I also had the pleasure to explore some interesting questions from the "how do I architect an asyncio-based application" domain, specifically "can I mix asyncio/non-asyncio code?"

The short answer is that yes, but you probably shouldn't. It's better to make the choice early on and stick to either yee olde blocking, synchronous code or make everything non-blocking and dependent on the event loop.

A high-level explanation might go like this: mixing sync/async code makes it easy to introduce subtle errors into your code. Whenever a sync portion is executing, the event loop is paused. The sync parts of your code have no way to yield control over to the event loop short of calling loop.run_until_complete. In simple scenarios this might be enough and this could be a decent way to increase IO performance in key areas of your application. For example:

# normal, somewhat contrived, sync code
def authenticate_users(users):
    loop = asyncio.get_event_loop()
    # yield control to even loop to execute a list of tasks asynchronously
    results = loop.run_until_complete(_async_auth(user) for url in users)
    # back to sync mode of operation

In an ideal world this works alright. But let's assume that this is a long-running application and it receives a termination signal from the OS. It would be great if we could run a bit of cleanup code before exiting to close any sessions or (to stick to the above example) de-authenticate the users. Python provides two ways to handle these signals: the signal module and the AbstractEventLoop.addsignalhandler. Using these, we can define a function that will run our cleanup code:

def cleanup():
    tasks = asyncio.Task.all_tasks()
    for task in tasks:
        task.cancel()
    some_async_session.close()
    # further cleanup

This looks all good until you realize that when this function is called, it interrupts the event loop. The event loop isn't stopped, it's actually full of unfinished tasks and its in the "running" state. Canceling tasks requires the event loop to run through the tasks. We can't do that at this point because calling run_until_complete or run_forever results in RuntimeError('Event loop is running.'). Calling loop.close results in RuntimeError("Cannot close a running event loop"). Ok, what about loop.stop? If I understand the underlying code correctly, loop.stop sets a variable that tells the loop to stop at the nearest occasion - but since our single-threaded app is executing the cleanup code, the event loop will never reach that Utopian land of "stopped" because the loop isn't running.

In order to tackle this problem, we can perform the cleanup operations synchronously and let the event loop crash, which usually produces some warnings but is fine. This entails getting a hold of "opened" resources (users in this case) and making them available in the cleanup function. If this is the only instance of using asyncio in the application then this is simple. But the more of them, the harder it will become because the cleanup function can be called at any moment and so it has to have access to a large set of resources.

The best way forward is to go all-in on asyncio in the application. We can use AbstractEventLoop.add_signal_handler to register a cleanup callback that will schedule a coroutine that cleans up resources and cancels all tasks except itself. Once the cleanup is done, the coroutine returns, and, being the only task left in the event loop, causes the event loop to shut down, exiting the application. We can also get fancy and raise an exception (eg. _AbnormalTermination) that will set the applications exit code to a non-zero value:

# registered with loop.add_signal_handler
def cleanup():
    asyncio.ensure_future(_async_cleanup())

async def _async_cleanup():
    tasks = [task for task in asyncio.Task.all_tasks() if task not asyncio.Task.current_task()]
    for task in tasks:
        task.cancel()
    # further cleanup
    raise AbnormalTermination()

try:
    loop.run_until_complete(application.run())
except KeyboardInterrupt:
    print('Application shutdown by user')
except AbnormalTermination:
    print('Application terminated!')
    raise SystemExit(1)

The key takeaway here is that you should decide whether your project will use asyncio early on, which will massively help you make sound architectural decisions later on.