Atexit for asyncio

Ideas
1

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
2

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

1 Like
3

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.

4

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.

5

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
6

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.

7

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.

8

Just FYI, @minrk is the JupyterHub maintainer.

2 Likes
9

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.

3 Likes
10

Thanks @minrk, I just came across a need for this and the asyncio-atexit package looks perfect for my use case.

However it would be helpful if something like this was included in the standard library to mimic the already present atexit functionality. My case is similar to @minrk’s with closing asynchronous connection pools / sessions at event loop exit. This is especially tricky to manage in test scenarios that don’t always mimic the real world with many new event loops being created between tests for isolation purposes.

11

I am sorry for dropping the ball here. Reading the whole thing again I think this would be a fairly simple addition to asyncio and we can easily get it into 3.14, as long as someone here (@minrk or @jrocketdev? Or another reader interested in using the feature) could do the work.

We need an issue in the python GitHub repo describing the proposed new feature (it should stand on its own but may reference this thread so history is linked), then a PR introducing the code (it can be much simpler than the asyncio-atexit package because it doesn’t need to monkey-patch).

The PR will also need to include docs and tests.

We should double check that the new feature works well with asyncio.run() as well as with loops started in other ways. It should work when the loop is interrupted with ^C or when it exits by someone calling .close(), or whatever other idioms are typically used by asyncio users in production.

3 Likes