Atexit for asyncio

It’s come up a few times that I want something like atexit behavior, but that should happen when the current asyncio event loop closes, rather than at interpreter exit. One big reason is because it may involve coroutines with references to the current loop (closing aiohttp Sessions, what have you).

The following patch seems to work for my purposes, but I don’t know what cases it might not handle:

def asyncio_atexit(callback):
    """
    Like atexit, but run when the asyncio loop is closing,
    rather than process cleanup.
    """

    loop = asyncio.get_running_loop()
    _patch_asyncio_atexit(loop)
    loop._asyncio_atexit_callbacks.append(callback)

async def _run_asyncio_atexits(loop):
    """Run asyncio atexit callbacks

    This runs in EventLoop.close() prior to actually closing the loop
    """
    for callback in loop._asyncio_atexit_callbacks:
        try:
            f = callback()
            if inspect.isawaitable(f):
                await f
        except Exception as e:
            print(f"Unhandled exception in asyncio atexit callback {callback}: {e}", file=sys.stderr)

def _asyncio_atexit_close(loop):
    """Patched EventLoop.close method to run atexit callbacks

    prior to the unpatched close method.
    """
    if loop._asyncio_atexit_callbacks:
        loop.run_until_complete(loop._run_asyncio_atexits())
    loop._asyncio_atexit_callbacks = None
    return loop._orig_close()

def _patch_asyncio_atexit(loop):
    """Patch an asyncio.EventLoop to support atexit callbacks"""
    if hasattr(loop, "_run_asyncio_atexits"):
        return

    loop._run_asyncio_atexits = partial(_run_asyncio_atexits, loop)
    loop._asyncio_atexit_callbacks = []
    loop._orig_close = loop.close
    loop.close = partial(_asyncio_atexit_close, loop)

Is something like this desirable? If not, what’s the recommended way for dealing with cleanup operations for “long running” tasks that should run for the lifetime of the loop, but might need some explicit cleanup at the end?

1 Like

I would create a custom event loop policy that installs a subclass of BaseEventLoop and override close() in the latter.

1 Like

Good idea, thanks! I guess my question is whether this is a common enough need that it warrants implementation in BaseEventLoop so it can be more widely adopted.

Setting a policy with a custom implementation works for applications, but anything solved at the policy level doesn’t help with libraries that want to be compatible with more than a couple asyncio-based applications controlled by the same team, since libraries have no way to appropriately influence the policy. We have the same problem with ProactorEventLoop’s incomplete API.

I don’t know what real world use case you have or how common it would be. That’s up to you to research and report.

Another option is passing ‘finalizer’ trampoline finction to asyncio,run():

async def finalize(coro):
    try:
        return await coro
    finally:
        for finalizer in finalizers:
            await finalizer()

asyncio.run(finalize(main())
2 Likes

Fair enough. I would think every case atexit is used for would apply, as folks move more things to async implementations. asyncio.run() in an atexit callback may work, as long as the cleanup coroutines don’t have any references to the current loop (e.g. closing connections attached to the current loop). Here is an example on StackOverflow.

Specifically, I’m working on this where KubeSpawner has recently adopted kubernetes_asyncio and thereby aiohttp. Many instances are created and destroyed during the lifetime of the application. For efficiency, they use a shared Client object, created when the first instance is initiated, which ultimately means a long-lived aiohttp.ClientSession that needs to be closed explicitly to avoid warnings in __del__ (related issue: __del__ cannot call cleanup methods if they are async). There are also long-running ‘reflector’ tasks using this client that run until the loop is closed. There is no hook provided (by asyncio or JupyterHub) to signal that the loop is shutting down, which is where I need to register the call to await client.close(). atexit would work fine for my purposes if client.close were not a coroutine, though tying it to the loop’s lifecycle would be even better.

I could address this by proposing a pytest-like addfinalizer API to JupyterHub, but this seems like a generic issue that all the problems solved by ‘atexit’ don’t have an async-compatible counterpart.

I like this a lot, though it’s not available for my current task. The package I’m working on is a plugin and not imported before the event loop is running, so it doesn’t have access to how asyncio.run is called.

I recommend taking this up with the JupyterHub maintainers first. If they say it needs additions to asyncio let them file a bug report here. For now, this seems premature.

Just FYI, @minrk is the JupyterHub maintainer.

2 Likes

Yup, I am a maintainer on both sides, so it is within my power to make cleanup-at-exit a JupyterHub-specific feature. That doesn’t solve the issue for more than one case, of course. It seems to me like a very generic feature to be able to register cleanup for long-running async tasks, though (as generic as atexit itself for folks using asyncio), which is why I brought it up here in case anyone else thinks it’s a good idea, or had their own solutions that can be accessible from an asyncio library rather than application.

FWIW, I packaged up my prototype as asyncio-atexit, since I think it’ll work in a wider variety of cases than modifications to event loop invocation.

2 Likes