[discussion] What to do about cancel scopes / nurseries inside generators?

[njsmith]

So here's a surprising thing that can happen in trio right now: if you have a generator (or async generator, doesn't matter) that yields inside a cancel scope, then the cancel scope remains in effect until you resume the generator and exit it.

For context managers, this is actually what you want, and we make use of this in lots of places, e.g.:

@contextmanager
def fail_at(deadline):
    with move_on_at(deadline) as scope:
        yield scope
    ...

with fail_at(deadline):
    XX

So here we have: first fail_at runs for a bit, and it starts a cancel scope, and then it yields. And then the XX code runs with the cancel scope in effect, even though failt_at is not on the stack. And then fail_at is resumed, and it exits the cancel scope. A little confusing when you think about the details, but basically this is safe and useful and ultimately does what you'd expect.

However, in general this can cause strange behavior, and also crashes:

import trio

def weird_generator():
    with trio.open_cancel_scope() as cs:
        yield cs

async def main():
    # hold onto a reference to prevent the gc from making this even more confusing:
    gen = weird_generator()
    for cs in gen:
        break
    cs.cancel()
    await trio.sleep(1)

trio.run(main)

Notice that even after we leave the loop and abandon the generator, the cancel scope is still in effect... and then the program crashes because trio detects that the cancel scope stack has become inconsistent (it tries to pop a scope that isn't at the top of the stack).

This all also applies to using nurseries inside async generators, because nurseries have an implicit cancel scope.

There are also some more generic issues with cleaning up (async) generators that call trio APIs after yielding, but that's another issue (specifically, issue #265). The cancel scope leakage causes problems even while the generator is alive, and even if we had PEP 533 then this would still be an issue.

Really the ideal solution here would be to error out if someone tries to yield through a cancel scope, IFF they're not implementing a context manager.

I don't think there's any way to do this though without PEP 521, and even then we'd have to figure out somehow whether this particular generator is implementing a context manager. (I guess some annoying thread-local thing could work -- e.g. make with trio.this_is_a_context_manager: ... which sets some internal flag that disables the PEP 521-based checks.)

I don't currently have any other ideas how to improve this, except to put some prominent warning in the documentation.

[njsmith]

Some discussion on python-ideas (scroll down to point 9): https://mail.python.org/pipermail/python-ideas/2017-August/046736.html

[njsmith]

Notes:

• PEP 568 provides one possible path for making cancel scopes in (async) generators "just work"; discussing this is a todo for python 3.8.

• There's some discussion that's relevant to this issue starting at Implement a strategy for handling (async) generator cleanup #265 (comment)

[oremanj]

One idea for nurseries-in-generators is to provide an ergonomic way to get a context-managed generator. I wrote an example implementation up at https://gist.github.com/oremanj/772241665edea3ceb3183a4924aecc08 (it was originally intended to be a quick-and-dirty demo, but rather grew a life of its own...) I think something like this would likely be useful in a trio utilities library. It doesn't require any special support from trio or async_generator and is rather large/magical so I'm not sure it deserves to make it into trio itself, unless trio wants to use it.

There are still notable limitations; for example, in the azip() example, if one of the inner slow_counters reaches its deadline and is cancelled, the whole azip() will be cancelled, because we had to enter the slow_counter contexts to get async iterables that could be passed to the azip() context. But if the limitations are acceptable, cancel scopes and nurseries that are attached to a generator in this fashion should be guaranteed to nest correctly / not mess up the cancel stack.

[njsmith]

Thinking about this a bit more...

PEP 568, if implemented, would make it possible to assign a reasonable semantics to cancel scopes with yield inside them. Basically, the scope would apply to the code that's statically inside the scope, but not to anything outside the generator; if the scope is cancelled/timeout expires while the generator is suspended, then that just means that when the generator is resumed, it'll be in a cancelled state and the next checkpoint will raise Cancelled. (Or if it exits the cancel scope without executing a checkpoint, that's fine too.)

PEP 568 would not make it possible to assign a reasonable semantics to a nursery with a yield inside it – that's intrinsically impossible if we want to keep the guarantee that exceptions in background tasks are propagated promptly, because if a background task crashes while the generator-holding-the-nursery is suspended, what can you do? There's no live stack frame to propagate it into. If we were willing to relax the guarantee to just, exceptions are propagated eventually, and PEP 533 were implemented, then we could allow yield inside nurseries, but neither of those seems likely.

So, maybe we want to we forbid yielding out of nurseries. How can we do that? One idea would be to have PEP 568, use it for cancel scopes, and then for nurseries add an additional feature: the ability to somehow "pin" the Context that's on top of the stack (the one that writes go to), so that if a yield would take us out of that context, then it fails (raises an exception instead of yielding). Then regular cancel scopes would stash their state in the context, and nurseries would additionally pin the context. And inside an @asynccontextmanager function you could still yield out of a nursery, because @asynccontextmanager functions – unlike other async generators – share their context with their caller, so yield there doesn't touch the context stack.

One more thing to consider: what if we decide to disallow yield inside all cancel scopes, not just nurseries? Even the correct semantics for a yield inside a cancel scope are a little confusing; if the timeout expires while the generator is suspended, then nothing happens, until later it's resumed and then boom. Users could conceivably be expecting a more complex semantics, where the timeout clock stops ticking during the yield. Or even if they aren't, then it's pretty weird for the semantics of the generator to depend on how long it was suspended for – that violates encapsulation, and probably isn't very useful.

If that's what we want to do, then we don't need to tie to it PEP 568-style nested contexts. The reason why tying them together is attractive is that PEP 568-style nested contexts need to have a special case to handle @contextmanager and @asynccontextmanager, and the no-yield-in-nursery rules need the exact same special case, so maybe we can combine the special cases. But we can also imagine adding the no-yield-in-nursery feature more directly, and decoupling the whole thing from PEP 568.

I guess what we'd want is: a way to set a flag on a caller's frame that makes it an error to attempt to yield out of that frame (but await is fine). And by "caller" I really mean "walk up the frame stack until we find the first frame that is not an @(async)contextmanager, or quit early if we encounter a frame that's not an (async) generator". On CPython this would actually be ... pretty easy to implement, and efficient. (Except that we might need some tweaks for the bytecode eval loop to distinguish between the different kinds of yield/await.) The PyPy folks might hate it; should check with them. Generators are a bit special AFAIK (you have to materialize their frames for suspension to work), so maybe it's not so bad?

Summary

Our options are:

• Allow yield in cancel scopes and nurseries
  • Requires: PEP 568, PEP 533, and relaxing the exception guarantees that nurseries currently provide
• Allow yield in cancel scopes, but not nurseries
  • Requires: PEP 568, plus some kind of "disallow yielding out of the current context" feature
• Disallow yield in both cancel scopes and nurseries
  • Requires: some kind of "disallow yielding the current frame" extension

PEP 533 is still how Python ought to work, but may not be the right hill to die on...

[njsmith]

walk up the frame stack until we find the first frame that is not an @(async)contextmanager, or quit early if we encounter a frame that's not an (async) generator

Actually, that's not quite right...

class CM1:
    def __enter__(self):
        self._cm2 = CM2()
        self._cm2.__enter__()

class CM2:
    def __enter__(self):
        disallow_yield()

with CM1():
    yield  # should be an error

So we really would want something like the PEP 568 semantics... like a stack of bools (or rather ints, to allow inc/dec, so disallow_yields can be nested), where the top of the stack tracks the allowed-to-yield state, entering an generator pushes to the stack, leaving pops, the push/pop is disabled on context manager generators, and the stack is automatically saved/restored when we suspend/resume coroutines, just like the exception stack. If we decide we want PEP 568 anyway, this could even be a ContextVar. Though the "pinning" idea has the advantage that copy_context wouldn't copy the disallow-yield state.

Eh, we have other stacks in the interpreter state, like the exception stack, it wouldn't be too weird to add another. And that way of formulating it wouldn't be a problem for PyPy.

[njsmith]

Example of a case where having the disallow-yield state be stored in the Context would be bad (especially if it's preserved across copy_context): the kluges pytest-trio uses to share a Context between fixtures and the test. See: python-trio/pytest-trio#59

[njsmith]

Frightening thought: we did a tremendous amount of work in python-trio/pytest-trio#55 to allow pytest-trio fixtures to yield inside nursery blocks. That work in particular is OK, because whatever hack we use to make @contextmanager work can also be applied to @pytest.fixture. (And it'll be a hack, but probably OK; the @pytest.fixture code does have some disturbingly detailed knowledge of why exactly yield is not OK and works around it, so I guess if other libraries started using disallow-yield for other reasons that could be a problem, but probably it's fine.) That's not the frightening thought.

The frightening thought is: well, why did we do all that work? It was because we decided it's not reasonable to expect people to know whether a random async context manager like async with open_websocket actually has a nursery inside it – that should be abstracted away, an internal implementation detail. But...... it can't be, because if you can't use nurseries inside random generators, then you can't use async with open_websocket inside random generators, and that means "does this context manager have a nursery inside it?" is actually a publicly visible feature of your API, and adding a nursery to an existing context manager may be a backwards-compatibility-breaking change :-( :-( :-(