Trio: Give a better error for improper combinations of AsyncExitStack with nurseries

[vxgmichel]

Hi all, I ran into a weird issue today. Consider the following program:

import trio
from contextlib import AsyncExitStack


async def start_context(stack):
    nursery = await stack.enter_async_context(trio.open_nursery())
    nursery.start_soon(trio.sleep, 1)


async def main():
    async with AsyncExitStack() as stack:
        async with trio.open_nursery() as nursery:
            nursery.start_soon(start_context, stack)


if __name__ == "__main__":
    trio.run(main)

It opens an AsyncExitStack and uses it to enter a nursery from another task. This produces the following traceback:

Traceback (most recent call last):
  File "/home/vinmic/miniconda/lib/python3.7/site-packages/trio/_core/_run.py", line 1769, in run
    run_impl(runner, async_fn, args)
  File "/home/vinmic/miniconda/lib/python3.7/site-packages/trio/_core/_run.py", line 1918, in run_impl
    runner.task_exited(task, final_outcome)
  File "/home/vinmic/miniconda/lib/python3.7/site-packages/trio/_core/_run.py", line 1391, in task_exited
    self.tasks.remove(task)
KeyError: <Task '__main__.start_context' at 0x7f7883818240>

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_trio.py", line 17, in <module>
    trio.run(main)
  File "/home/vinmic/miniconda/lib/python3.7/site-packages/trio/_core/_run.py", line 1775, in run
    ) from exc
trio.TrioInternalError: internal error in Trio - please file a bug!
Exception ignored in: <function Nursery.__del__ at 0x7f7883abca60>
Traceback (most recent call last):
  File "/home/vinmic/miniconda/lib/python3.7/site-packages/trio/_core/_run.py", line 964, in __del__
AssertionError:

I suspect this issue to be caused by the fact the the nursery is not entered and exited in the same task. As far as I understand this is a trio limitation. In this case, a nicer exception could really help :)

[oremanj]

It's not so much about tasks as about stack discipline -- you're doing

• enter outer nursery
• enter async-exit-stack nursery
• exit outer nursery
• exit async-exit-stack nursery

but Trio expects you to exit nurseries in the opposite order you entered them, like you would if you'd used async with statements. The error here should definitely be better; we cleaned up the error handling for this sort of misuse with cancel scopes, but haven't gotten to nurseries yet.

If I understand your goal correctly, you could do something like

async def run_context(*, task_status=trio.TASK_STATUS_IGNORED):
    async with trio.open_nursery() as nursery:
        task_status.started(nursery)
        nursery.start_soon(trio.sleep, 1)

async def main():
    async with trio.open_nursery() as outer_nursery:
        inner_nursery = await outer_nursery.start(run_context)
        # and then spawn stuff into inner_nursery as needed

The sleep(1) task is bit of a hack; there was some discussion of less sleepy alternatives in #304.

[vxgmichel]

It's not so much about tasks as about stack discipline

I agree that this is sloppy code (and probably conceptually wrong), but I don't think it's only about the incorrect nesting of the nurseries. For instance, the following example produces a similar trace even though the nurseries are correctly nested:

async def start_context(stack, task_status):
    nursery = await stack.enter_async_context(trio.open_nursery())
    task_status.started(nursery)


async def main():
    async with trio.open_nursery() as nursery:
        async with AsyncExitStack() as stack:
            context = await nursery.start(start_context, stack)
            context.start_soon(trio.sleep, 1)

Here's a more realistic example, showing how one might run into this issue without really noticing the sloppy logic. Obviously, the culprit here is:

        async def get_service(self, n):
            if n not in self.cache:
                self.cache[n] = await stack.enter_async_context(some_service())
            return self.cache[n]

That's when I realized that AsyncExitStack is a bit dangerous and that we're usually better off using a nursery to manage async contexts. Something along the lines of:

        async def get_service(self, n):

            async def target(task_status):
                async with some_service() as service:
                    task_status.started(service)
                    await trio.sleep_forever()

            if n not in self.cache:
                self.cache[n] = await nursery.start(target)

            return self.cache[n]

(which is very similar to the example you provided).

I don't know if this is a good reason to discourage the use of AsyncExitStack altogether, I just thought I'd report it to show that those situations can actually occur in real-life code :)

[smurfix]

we're usually better off using a nursery to manage async contexts
I don't know if this is a good reason to discourage the use of AsyncExitStack altogether

AsyncExitStack is very useful IMHO.

The problem isn't using a nursery vs. an exit stack to manage exits, the problem is that you're trying to use an exit stack to manage a nursery. That's possible, no problem, but you need to create the nursery within an @asynccontextmanager-decorated function.

This code works:

import trio
from contextlib import AsyncExitStack, asynccontextmanager

@asynccontextmanager
async def start_context():
    async with trio.open_nursery() as n:
        yield n

async def main():
    async with AsyncExitStack() as stack:
        context = await stack.enter_async_context(start_context())
        context.start_soon(trio.sleep, 1)

trio.run(main)

[vxgmichel]

@smurfix

This code works

Oh yes, it works as long as enter_async_context runs in the same task as AsyncExitStack. But running enter_async_context in another task reproduces the issue I initially reported:

async def main():
    async with trio.open_nursery() as nursery:
        async with AsyncExitStack() as stack:

            async def target():
                context = await stack.enter_async_context(start_context())
                context.start_soon(trio.sleep, 1)

            nursery.start_soon(target)
            await trio.sleep(.1)

AsyncExitStack is very useful IMHO.

Maybe trio could provide its own enhanced version of AsyncExitStack, making sure that all its sensitive methods run in the same task in which the stack has been entered?

[belm0]

context = await stack.enter_async_context(start_context())

I don't have any problem passing trio.open_nursery() to enter_async_context directly.

Yes, I learned the hard way that you can't play games with nursery enter and exits. Exits need to be in stack order and from the same task. You can usually structure your concurrency around this limitation.

[smurfix]

But running enter_async_context in another task reproduces the issue I initially reported:

The problem isn't that you're doing it in another task. The problem is that you're dismantling things in some random order, i.e. not reversing the order you've set them up in.

Don't do that.

Nurseries go splat in a somewhat spectacular way when you do it, because they intrinsically depend on that order, but you can construct the same problem with any other pair of resources where one depends on the other.

[njsmith]

Yeah, there's two separate constraints: you have to exit nurseries in the same order you entered them, and you have to exit nurseries from the same task where you entered them.

Basically nurseries are designed around async with statements, and they only support patterns that can be written using actual async with statements. Since this is python, we can't stop you from doing tricky stuff like calling __aenter__ directly or using AsyncExitStack in creative ways. And if the end result of this trickery is equivalent to something you could write with plain async with statements, then that's fine (though maybe not very useful). But if you're trying to use the trickery to let you write things that couldn't be written using async with, then that's not supported.

Maybe trio could provide its own enhanced version of AsyncExitStack, making sure that all its sensitive methods run in the same task in which the stack has been entered?

I mean, that would be possible to do. It would also be possible to get rid of nurseries altogether, and just let you spawn tasks whenever you wanted. It would be even easier than implementing nurseries, in fact :-). But the reason trio has nurseries is because we think that the constraint to use async with actually makes for better code. And all the cases I know of where you would want to use a AsyncExitStack-with-special-nursery-support and it wouldn't violate the structured concurrency constraints, there are already other better ways to write it. So I'm not seeing a compelling case for adding this.

[vxgmichel]

@njsmith

Yeah, there's two separate constraints: you have to exit nurseries in the same order you entered them, and you have to exit nurseries from the same task where you entered them.

Basically nurseries are designed around async with statements, and they only support patterns that can be written using actual async with statements.

Alright, that makes sense!

I mean, that would be possible to do. It would also be possible to get rid of nurseries altogether, and just let you spawn tasks whenever you wanted.

Oh sorry I wrote this part to quickly, that's not what I meant. I just wondered it it was possible to check the constraints you mentioned in advance in order to raise an exception earlier, since it is easy to use AsyncExitStack incorrectly. For instance, I just realized this code also fails:

    async with AsyncExitStack() as stack:
        async with trio.open_nursery() as nursery:
           service = await stack.enter_async_context(
               some_service_that_might_run_a_nursery())

This might seem obvious with small examples, but it very easy to miss in larger code base. In our case, we used AsyncExitStack incorrectly for months until a nursery was added to some apparently unrelated async context manager in some other part of the code base. Just like in the above example, everything works fine as long as some_service_that_might_run_a_nursery does not do so.

It would be even easier than implementing nurseries, in fact :-). But the reason trio has nurseries is because we think that the constraint to use async with actually makes for better code.

And that's much appreciated by the way :)

[njsmith]

Oh sorry I wrote this part to quickly, that's not what I meant. I just wondered it it was possible to check the constraints you mentioned in advance in order to raise an exception earlier, since it is easy to use AsyncExitStack incorrectly.

Hmm, that does sound nice, but I'm not sure what constraints we could enforce. In your case, it sounds like the code was actually correct until that second nursery got added? Do you have a heuristic in mind for how we could detect that your usage was risky, before it actually became a problem?

[vxgmichel]

@njsmith

Hmm, that does sound nice, but I'm not sure what constraints we could enforce. In your case, it sounds like the code was actually correct until that second nursery got added? Do you have a heuristic in mind for how we could detect that your usage was risky, before it actually became a problem?

This is roughly what I had in mind, although I might be missing something:

import trio
import contextlib


class AsyncExitStack(contextlib.AsyncExitStack):
    async def __aenter__(self):
        self._task = trio.hazmat.current_task()
        self._current_scope = self._task._cancel_status._scope
        return await super().__aenter__()

    async def enter_async_context(self, cm):
        if self._task != trio.hazmat.current_task():
            raise RuntimeError(
                "Async context must be entered in the same task as the stack"
            )
        if self._current_scope != self._task._cancel_status._scope:
            raise RuntimeError(
                "Async context must be entered in the scope of the previous context"
            )
        try:
            return await super().enter_async_context(cm)
        finally:
            self._current_scope = self._task._cancel_status._scope

The enter_async_context method basically check the constraints you mentioned in this post.

It seems to work as expected:

import pytest

@pytest.mark.trio
async def test_async_exit_stack():

    # Proper nesting and single task
    async with trio.open_nursery() as nursery:
        async with AsyncExitStack() as stack:
            service1 = await stack.enter_async_context(trio.open_nursery())
            service2 = await stack.enter_async_context(trio.open_nursery())
            nursery.start_soon(trio.sleep, 0.1)
            service1.start_soon(trio.sleep, 0.1)
            service2.start_soon(trio.sleep, 0.1)
            await trio.sleep(0.01)

    # Invalid nesting
    async with AsyncExitStack() as stack:
        async with trio.open_nursery() as nursery:
            with pytest.raises(RuntimeError):
                await stack.enter_async_context(trio.open_nursery())

    # Multiple tasks
    async with trio.open_nursery() as nursery:
        async with AsyncExitStack() as stack:

            async def target():
                with pytest.raises(RuntimeError):
                    await stack.enter_async_context(trio.open_nursery())

            nursery.start_soon(target)
            await trio.sleep(0.01)

[belm0]

trio.TrioInternalError: internal error in Trio - please file a bug!

It's way more likely for an application to be doing something wrong with AsyncExitStack or explicit aenter/aexit calls than for there to be a Trio bug. The message here should at least reflect that. If it's possible to detect the issue slightly higher up in the stack and give a little more context in the error, that would be even better.

The error which comes up when calling yield from a nursery should get similar treatment. While it does suggest your application is probably doing something wrong, it should mention the nursery/generator limitation explicitly.