Nasty edge case in async generator shutdown: can't tell what order to close them in

On chat today, @SaschaSchlemmer / @SDesch was asking about the exception traceback you get when using context managers inside async generators without aclosing. I was a bit confused, because I didn't think you were supposed to get an exception traceback here... maybe a ResourceWarning at worst. But indeed, it turns out that this code does print an exception traceback, randomly some proportion of the time:

import sys
import trio
from contextlib import asynccontextmanager
import traceback

@asynccontextmanager
async def context():
    print('entering context')
    try:
        yield
    finally:
        print('leaving context', sys.exc_info())

async def agen():
    print("entering agen")
    try:
        async with context():
            for i in range(100):
                try:
                    yield i
                finally:
                    print("resuming agen", sys.exc_info())
    finally:
        print("exiting agen", sys.exc_info())

async def example():
    async for i in agen():
        break

for i in range(10):  # run a couple of times to see the warning/error
    print("-----")
    trio.run(example)

Here's a typical failing output:

entering agen
entering context
leaving context (<class 'GeneratorExit'>, GeneratorExit(), <traceback object at 0x7fc022500780>)
resuming agen (<class 'GeneratorExit'>, GeneratorExit(), <traceback object at 0x7fc022500800>)
exiting agen (<class 'RuntimeError'>, RuntimeError("generator didn't stop after athrow()"), <traceback object at 0x7fc022500a80>)
Exception ignored during finalization of async generator '__main__.agen' -- surround your use of the generator in 'async with aclosing(...):' to raise exceptions like this in the context where they're generated
Traceback (most recent call last):
  File "agen-weirdness.py", line 20, in agen
    yield i
GeneratorExit

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/njs/.user-python3.8/lib/python3.8/site-packages/trio/_core/_asyncgens.py", line 186, in _finalize_one
    await agen.aclose()
  File "agen-weirdness.py", line 22, in agen
    print("resuming agen", sys.exc_info())
  File "/usr/lib/python3.8/contextlib.py", line 190, in __aexit__
    raise RuntimeError("generator didn't stop after athrow()")
RuntimeError: generator didn't stop after athrow()

We can see that:

• agen enters context
• then before agen is resumed, we see leaving context

This is weird, because context is a context manager inside agen, so how can it be left before agen has resumed?

Answer: the context manager has an async generator object hidden inside it. When trio.run completes, it manually goes around and closes all outstanding async generator objects. They're stored in a WeakSet, so the order they're processed is random. So sometimes, Trio forces the hidden context manager async generator to close first, and then it closes agen. And while unwinding agen, it exits the context manager, and then contextlib gets very confused when it tries to throw the exception into context, but the context function has already terminated.

This is really nasty. We don't really have any reliable way to figure out which order async generators should be shut down. This might be like... a bug in PEP 525? (CC @1st1). Looking at asyncio's code, I think they have the same issue.

Another complication: I did try running the above code on asyncio, and it did not exhibit the problem. I was very confused for a while before I figured it out: for some reason, on trio the agen object remains alive until we're closing async generators, while on asyncio, it gets GC'ed, so asyncio closes it early, before the final forced shutdown.

I don't really understand how we could be holding the agen object alive here. I guess it must be something like... the frame object holds a reference to it while the iteration is happening. And then since this is the "main" task at the top of the task tree, trio internally holds onto it in Runner.main_task, and that holds the coroutine object alive, and for some reason the coroutine object holds the frame alive even though the coroutine has completed, and for some reason the frame holds the async iterator alive even though the for loop has exited? None of that makes much sense, but I can't think what else would be going on. I guess there's the StopIteration that gets raised when example exits, which also references example's frame, but AFAICT we drop our reference to the StopIteration immediately after reading out its value, so I don't think that's holding anything alive. This is probably a secondary bug: we should drop that frame immediately after the main task completes, rather than holding onto it.

If I'm right, this somewhat limits the impact of this bug: it would mean Sascha only hit it because they had exactly the right program structure and we has this secondary bug. But, the underlying problem definitely can arise in other cases -- if you have a non-refcounted GC (= pypy), then this will happen all the time; or, if your async generator genuinely is being kept alive somehow (perhaps through a reference loop, or simply because your main task terminated with an exception). And since an actual user actually got bitten by this and it horribly misled them about how async generators work, I feel like we should take it seriously.

At least in this case, we could mitigate it by processing async generator shutdowns in FIFO order -- that would mean we would always close agen first, before closing context. This is counterintuitive, since usually when unwinding things you use LIFO order, but the intuition here is that we want to shut down "users" before the things they're using, and usually users are created first. It's still a heuristic and not guaranteed to work in all cases, but maybe it covers 95%?

Also, I think on modern Python, weakref.WeakKeyDict is order-preserving? So maybe we just need to replace our WeakSet with that?

So tenative todo list:

• figure out how agen is being kept alive, and fix it
• track async generators in some kind of ordered data structure (not WeakSet)

Start with trio/_core/_asyncgens.py at _finalize_one, then inspect how Runner.main_task may retain the completed frame and async generator. Reproduce the provided example and investigate the WeakSet shutdown tracking. Done means understanding the retention behavior and establishing a reliable shutdown-order fix or documented limitation.