Why is __aexit__ not fully executed when it has await inside?

Question:

This is the simplified version of my code:

main is a coroutine which stops after the second iteration.
get_numbers is an async generator which yields numbers but within an async context manager.

import asyncio


class MyContextManager:
    async def __aenter__(self):
        print("Enter to the Context Manager...")
        return self

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        print(exc_type)
        print("Exit from the Context Manager...")
        await asyncio.sleep(1)
        print("This line is not executed")  # <-------------------
        await asyncio.sleep(1)


async def get_numbers():
    async with MyContextManager():
        for i in range(30):
            yield i


async def main():
    async for i in get_numbers():
        print(i)
        if i == 1:
            break


asyncio.run(main())

And the output is:

Enter to the Context Manager...
0
1
<class 'asyncio.exceptions.CancelledError'>
Exit from the Context Manager...

I have two questions actually:

  1. From my understanding, AsyncIO schedules a Task to be called soon in the next cycle of the event loop and gives __aexit__ a chance to execute. But the line print("This line is not executed") is not executed. Why is that? Is it correct to assume that if we have an await statement inside the __aexit__, the code after that line is not going to execute at all and we shouldn’t rely on that for cleaning?

  1. Output of the help() on async generators shows that:
 |  aclose(...)
 |      aclose() -> raise GeneratorExit inside generator.

so why I get <class 'asyncio.exceptions.CancelledError'> exception inside the __aexit__ ?

* I’m using Python 3.10.4

Asked By: S.B

||

Answers:

To answer the first question:

Is it correct to assume that if we have an await statement inside the
__aexit__, the codes after that line is not going to execute at all?

I would say no it’s not always the case. As long as the main has enough time and can again pass the control back to event loop, the code inside the __aexit__ may execute. I tried this:

async def main():
    async for i in get_numbers():
        print(i)
        if i == 1:
            break
    await asyncio.sleep(4)   #  <---- New

.run() only cares about the coroutine which is passed to it and runs that to the end, not the other coroutines including __aexit__. So if it doesn’t have enough time or doesn’t pass the control to event loop, I can’t rely on the next lines after the first await statement.


Additional information that may help:

Here in base_events.py/run_forever method(Which invoked by .run()), I found that self._asyncgen_finalizer_hook is passed to the sys.set_asyncgen_hooks. The body of the _asyncgen_finalizer_hook is :

    def _asyncgen_finalizer_hook(self, agen):
        self._asyncgens.discard(agen)
        if not self.is_closed():
            self.call_soon_threadsafe(self.create_task, agen.aclose())

But the implementation of call_soon_threadsafe is empty.

I’ll clean up this answer and remove these guesses later.

Answered By: S.B

I’m not sure what’s happening, but posting what I found in case it proves useful to others who decide to investigate. The output changes to expected when we store reference to get_numbers() outside main(). I would say that it seems get_numbers() is garbage collected to early, but disabling gc doesn’t help, so my guess may be wrong.

import asyncio


test = None


class MyContextManager:
    async def __aenter__(self):
        print("Enter to the Context Manager...")
        return self

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        print(exc_type)
        print("Exit from the Context Manager...")
        await asyncio.sleep(1)
        print("This line is not executed")  # <-- Executed now
        await asyncio.sleep(1)


async def get_numbers():
    async with MyContextManager():
        for i in range(30):
            yield i


async def main():
    global test
    test = get_numbers()

    async for i in test:
        print(i)
        if i == 1:
            break


asyncio.run(main())
Answered By: Mikhail Gerasimov

the answer is simple: the interpreter will continue executing __aexit__ after one second, but the main function is finished and there is no pointer to the context manager.

first obvious solution which is mentioned by yourself is to wait long enough after the main function:

async def main():
    async for i in get_numbers():
        print(i)
        if i == 1:
            break
    await asyncio.sleep(4)   #  <---- New

another way is to use try/finally:

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        try:
            pass
            print(exc_type)
            print("Exit from the Context Manager...")
            await asyncio.sleep(1)
        finally:
            print("This line is not executed")  # <-------------------
Answered By: Kasra

This is not specific to __aexit__ but to all async code: When an event loop shuts down it must decide between cancelling remaining tasks or preserving them. In the interest of cleanup, most frameworks prefer cancellation instead of relying on the programmer to clean up preserved tasks later on.

This kind of shutdown cleanup is a separate mechanism from the graceful unrolling of functions, contexts and similar on the call stack during normal execution. A context manager that must also clean up during cancellation must be specifically prepared for it. Still, in many cases it is fine not to be prepared for this since many resources fail safe by themselves.


In contemporary event loop frameworks there are usually three levels of cleanup:

  • Unrolling: The __aexit__ is called when the scope ends and might receive an exception that triggered the unrolling as an argument. Cleanup is expected to be delayed as long as necessary. This is comparable to __exit__ running synchronous code.
  • Cancellation: The __aexit__ may receive a CancelledError1 as an argument or as an exception on any await/async for/async with. Cleanup may delay this, but is expected to proceed as fast as possible. This is comparable to KeyboardInterrupt cancelling synchronous code.
  • Closing: The __aexit__ may receive a GeneratorExit as an argument or as an exception on any await/async for/async with. Cleanup must proceed as fast as possible. This is comparable to GeneratorExit closing a synchronous generator.

To handle cancellation/closing, any async code – be it in __aexit__ or elsewhere – must expect to handle CancelledError or GeneratorExit. While the former may be delayed or suppressed, the latter should be dealt with immediately and synchronously2.

    async def __aexit__(self, exc_type, exc_value, exc_tb):
        print("Exit from the Context Manager...")
        try:
            await asyncio.sleep(1)  # an exception may arrive here
        except GeneratorExit:
            print("Exit stage left NOW")
            raise
        except asyncio.CancelledError:
            print("Got cancelled, just cleaning up a few things...")
            await asyncio.sleep(0.5)
            raise
        else:
            print("Nothing to see here, taking my time on the way out")
            await asyncio.sleep(1)

Note: It is generally not possible to exhaustively handle these cases. Different forms of cleanup may interrupt one another, such as unrolling being cancelled and then closed. Handling cleanup is only possible on a best effort basis; robust cleanup is achieved by fail safety, for example via transactions, instead of explicit cleanup.


Cleanup of asynchronous generators in specific is a tricky case since they can be cleaned up by all cases at once: Unrolling as the generator finishes, cancellation as the owning task is destroyed or closing as the generator is garbage collected. The order at which the cleanup signals arrive is implementation dependent.

The proper way to address this is not to rely on implicit cleanup in the first place. Instead, every coroutine should make sure that all its child resources are closed before the parent exits. Notably, an async generator may hold resources and needs closing.

async def main():
    # create a generator that might need cleanup
    async_iter = get_numbers()
    async for i in async_iter:
        print(i)
        if i == 1:
            break
    # wait for generator clean up before exiting
    await async_iter.aclose()

In recent versions, this pattern is codified via the aclosing context manager.

from contextlib import aclosing

async def main():
    # create a generator and prepare for its cleanup
    async with aclosing(get_numbers()) as async_iter:
        async for i in async_iter:
            print(i)
            if i == 1:
                break

1The name and/or identity of this exception may vary.

2While it is possible to await asynchronous things during GeneratorExit, they may not yield to the event loop. A synchronous interface is advantageous to enforce this.

Answered By: MisterMiyagi