Implement a strategy for handling (async) generator cleanup

[njsmith]

There's a problem with generators and async generators: they don't necessarily iterate to the end. Sometimes they get suspended in the middle, and then garbage collected. When this happens, their __del__ method throws in a GeneratorExit exception, so that finally blocks and __exit__ methods get a chance to run. BUT, __del__ is very weird: it can be run in the middle of any arbitrary python bytecode instruction. It's like a signal handler but even more so.

So for example, this code:

def generator_fn():
    try:
        yield
    finally:
        print(trio.current_task())

async def caller():
    for _ in generator_fn():
        break

Here the generator's finally block gets run at some totally arbitrary time, so it could in principle print any task, or crash because no task is currently active, or because it doesn't even run in the trio thread. And similar problems arise for pretty much any access to trio APIs. (Note that even without trio, the same thing happens for threading.current_thread(), accessing thread-local storage, etc. But it trio the use of this kind of implicit thread-bound state is much more common, since the run loop itself is accessed via thread-local storage.)

PEP 533 would be the real fix, but that seems to be stalled for now.

For regular generators, we just need to document this; I can't see what else we can do. There are some related issues in #264 that also need documentation; probably they should go in the same section of the manual.

For async generators, this problem is both better and worse. The thing about async generators is that the gc can't actually clean them up directly: it would like to do the same thing as it does for regular generators and have __del__ throw in GeneratorExit, but this doesn't really work, because for an async generator you should use athrow and let the coroutine runner iterate it, but __del__ doesn't have access to a coroutine runner. So PEP 525 defines some hooks that trio can set, to get (a) notification that an async generator has been started, and (b) notification when an async generator is ready for garbage collection.

The intention was that these hooks are used to let the event loop take care of doing gen.athrow(GeneratorExit) to emulate how regular generators work. But as we saw above, regular generators are already broken so... we don't want to emulate them. Being able to re-enter trio's run loop doesn't help, because we need a task context. And this is even more true for async generators, since the whole point of an async generator is that it can do async-y things, and if you want to await something you need to make sure you have a timeout set, which means you really really really need to make sure that this happens in the original task context where the user was expecting it to happen. Letting an await sneak outside of its with fail_after is definitely not OK.

So we aren't going to use these hooks in the way they were originally intended. But, until PEP 533 becomes a thing, maybe we can at least use these hooks for harm reduction.

Some options:

• Don't define the hooks at all. In this case I'm not 100% sure what Python does, but I think it does something like: throw in GeneratorExit and iterate once, and if it yields again (e.g. because there's an async with or a finally: await ...) then it gives up and prints an error message. This is basically like synchronous generators, but maybe a little less friendly. It's not awful, but maybe we can do better.

• Use the firstiter hook to forbid the use of async generators, unless we have some explicit evidence that someone is taking responsibility for cleaning them up. For example, we could make it so that our @acontextmanager marks its async generators as being cleaned up (though 3.7's contextlib.asynccontextmanager is going to be a problem...), and provide an aclosing context manager that does the same. This approach was pioneered by Curio, which goes a step further and introspects the generator's bytecode, and if it doesn't see any dangerous operations (yield or yield from inside an except, finally, with, or async with), then it allows those async generators as well. I really don't want to have this kind of bytecode introspection in trio. OTOH without it then async comprehensions like (2 * i async for i in aiter) become essentially unusable.

• Use the finalizer hook to catch when an async generator is cleaned up, and run it in a context that (a) is carefully set up so that any attempt to access trio state generates a clear error message, and (b) any attempt to yield causes a clear error / crashes the program (as opposed to just being printed to stderr, like regular __del__ does). (If doing this we'd probably also want a firstiter hook like asyncio's, so we can force all async generators to be finalized when the run loop exits. This is just a minor refinement, though.)

The good thing about the finalizer approach is that it has no false negatives: all async generators that can be run safely are allowed to run. The downside is that it has false positives: if an async generator is written poorly, then it's possible that it will be allowed to run and everything will look fine in tests, but then in production it blows up (e.g. because in tests it's always iterated to the end, and in production someone breaks out of a loop so the generator gets collected while only half-finished). The firstiter approach is the opposite: it only allows generators to run if it can be sure that they're safe in all circumstances.

(If PEP 533 is too much for 3.7, then maybe a smaller ask would be set_gen_hooks() that work the same way but for synchronous generators?)

[oremanj]

Here's an idea: what if we scope async generators to the task that first iterated them? That is:

• each task contains an _asyncgens attribute containing a weakref list (can be initially None to avoid having a bunch of empty sets for the tasks that don't use asyncgens)
• we install a firstiter hook that adds a weakref to the new async generator to the _asyncgens list on the current task
• when a child task of nursery exits, we start a task in the nursery that does the equivalent of await ag.aclose() for each of its still-alive asyncgens, in reverse order of firstiter

There's some subtlety here as far as treating the nested child sorta-task like a task for these purposes, and scoping asyncgens created there to the direct-parent nursery rather than to the lifetime of the overarching task that opened that nursery, but I think that's mostly details.

This is only a partial solution; it still doesn't deal correctly with asyncgens that create nurseries/cancel scopes.

async def agen():
    with trio.move_on_after(6):
        for i in range(100):
            yield i
            await trio.sleep(i * 0.1)

async def caller():
    ag = agen
    last = None
    while True:
        with trio.move_on_after(0.8) as scope:
            last = await ag.anext()
        await trio.sleep(1)
        if scope.cancelled_caught:
            break

From reading the discussion in #264, it seems like the desired behavior here would be to limit each turn of agen() to 0.8 seconds, and the whole thing (including time spent by caller() when agen() is suspended) to 6 seconds. That gives you a shorter-lived cancel scope (in caller()) more deeply in the cancel scope stack than a longer-lived one (in agen()). I guess context variables are supposed to solve this? I feel vaguely like it should be solvable without them, but can't currently think of how. I guess cancel scopes could detect that they were created inside an asyncgen (via... stack introspection? ugh) and mark themselves accordingly, and cancel scope __exit__ could know it might need to reorder such scopes and delivery of cancels would skip them if the asyncgen is suspended due to a yield. This gets really complicated though. Thoughts?

[oremanj]

Thinking about this more:

• The fundamental problem for tracking of cancel scopes and other contexts within an asyncgen is knowing whether it's currently suspended by a yield (to its async caller) or by a trap (to the event loop) yielded by some other async function it awaited. An asyncgen can't directly yield to the event loop, because "yield" in the body of an async function syntactically means yield to asyncgen caller. So any event loop trap must be behind an await.
• We know when asyncgens are created and destroyed, if we install the hooks appropriately. We can scope the lifetime of an asyncgen to the lifetime of the task that created it, as explained above. So, whenever a task traps to the event loop, we know what asyncgens might be running at the time of the trap.
• We can iterate over each of these asyncgens and check its ag_await member. If this is None, the asyncgen is suspended by a yield to its caller and we shouldn't consider its contexts to be in effect at the time of the trap. If it's not None, the asyncgen is suspended by awaiting someone else and we should consider its contexts to be in effect at the time of the trap.
• The missing piece is figuring out which contexts (pertaining to some task-local state like the cancel scope stack) "belong" to the asyncgen and which "belong" to the code that's calling it. I'll think about this some more.

[oremanj]

Expanding on that last bullet point: we want to distinguish

async def agen():
    with trio.move_on_after(1):
        yield 42
async def caller():
    async for val in agen:
        await trio.sleep(val)

from

async def agen():
    yield 42
async def caller():
    with trio.move_on_after(1):
        async for val in agen:
            await trio.sleep(val)

(the second should interrupt the sleep while the first should not). Luckily, the firstiter hook is called when the coroutine returned by asend() is constructed, before the send(None) is performed and thus before any of the async generator code is executed. So if we add asyncgens to Task._asyncgens in firstiter, we can have CancelScope.__enter__ (and any other relevantly context-y things) check each asyncgen's ag_running (works despite bpo-32526 since this is a sync call) and do whatever is necessary to remember that this cancel scope belongs to this asyncgen. I just checked and ag_running appears to be True from the perspective of an async function awaited from the generator; it just flips back to False when the call stack actually traps to the event loop. So I don't think the async-ness of open_nursery() will get in the way either.

I'm sure the devil will be in the details, but can you see any missing pieces in the overall idea?

None of this helps with regular non-async generators that create a cancel scope. But since there's not much point to making a cancel scope that you can't put any awaits in, I don't think there's any use for such a generator besides decorating it with @contextmanager, and we've established we like the existing behavior in that case.

[oremanj]

Even more wrinkles: asyncgens can iterate over other asyncgens! I don't have time to think through the implications of this right now; mostly leaving a note here for my future self. I can't immediately come up with any reasons this would invalidate the approach I'm proposing, but it probably does at least mean we need to be careful about the order in which we apply the asyncgen cancel scopes to the cancel scope stack.

Also, ag_await is None is not enough to know an asyncgen is suspended at a yield -- it could just be exhausted. The condition we want for is-currently-suspended-via-yield-to-caller is ag_await is None and ag_frame is not None.

[njsmith]

Random technical notes

I was definitely going to ask about bpo-32526; that's a pretty obscure issue, so nice catch :-).

It's also legal to create an async generator in one task and then iterate it in another task. Actually I've even used this in an example... in the live-coded happy eyeballs, I cheat a little bit, because I claim that it's doing the same thing as the twisted happy eyeballs code, but twisted's version actually supports incremental getaddrinfo. (Which isn't supported by any real OS, but never mind.) But it's easy to adapt the example to handle the full twisted semantics: suppose we have an incremental getaddrinfo that returns an async iterator:

async def fake_incremental_getaddrinfo(*args, **kwargs):
    for target in await getaddrinfo(*args, **kwargs):
        yield target

Then in the happy eyeballs code, each child task instead of doing target = targets[i] instead does:

try:
    target = await targets_aiter.__anext__()
except StopAsyncIteration:
    return

You also need to do a bit of fiddling with the failed_attempts variable because you don't know ahead of time how many attempts there will be, but never mind -- the point is that if anyone ever does ask about the twisted semantics, then it's also quite easy to do and only adds a few lines. But... it's also a realistic case where iterating a generator in different tasks would be useful.

I suppose one could keep a global list of all living async generators, and scan them whenever a cancel scope is created or a task traps, but this would be way too expensive. (It might even be too expensive if we could restrict to a single task's generators... these are pretty performance-sensitive operations.)

Another potential way to map cancel scopes to async generators would be to use the firstiter hook to keep a global set of living async generator frame objects (!), and then when a cancel scope is created, check whether any of them are on the stack. On CPython at least this would be pretty quick (because now it's O(stack depth), not O(total number of generators in the program)). On PyPy it might be OK, I'm not sure – forcing the JIT to materialize frame objects isn't great, but if we're not looking at the contents of the objects, and it's only for async function frames in practice, then maybe it's not too bad?

But let's step back a bit

The bigger question is what semantics we're even going for here. I think in this issue we should try to focus on just the GC issues, since they're simpler. (I added a link back to here from #264 though so the comments don't get forgotten.)

If you want to work on this, then probably the first thing that would be useful would be to figure out exactly what Python is doing right now, with no hooks defined :-). I wrote down a guess in the first post, but I haven't checked the details (which may require reading CPython's source and stuff).

There's also a simple thing we could do that might be better than what we do now, not sure: guarantee that all async generators will get await trio.aclose_forcefully(agen) called on them.