Why can't you await __aenter__? (Understanding trio's contracts around async context managers)

[yang]

Say we have context managers:

async with a(), b(): ...

We want to parallelize their enter/exit.

One approach: create a new context manager c that duplicates the logic from both, but ideally we don't have to do that—keep both a and b as reusable context managers, but also run them in parallel.

Instead, let's create a combinator:

async with par_with(a, b): ...

What's the right way to implement par_with?

First idea that came to mind:

@asynccontextmanager
async def par_with(a, b):
  ca, cb = a(), b()
  async with open_nursery() as n:
    n.start_soon(ca.__aenter__)
    n.start_soon(cb.__aenter__)
  yield
  async with open_nursery() as n:
    n.start_soon(cb.__aexit__)
    n.start_soon(ca.__aexit__)

But this causes nested nurseries to get canceled as soon as the enters finish. OK, so far this makes sense to me.

So we want to keep the outer nursery strictly surrounding both.

Another stab:

@asynccontextmanager
async def par_with(a, b):
    (ca, cb) = (a(), b())
    cn = trio.open_nursery()
    n = await cn.__aenter__()

    # This is so that we can wait for both
    done_enter = trio.Event()
    count = 0

    async def call_and_signal(f):
        nonlocal count
        n.start_soon(f)
        count += 1
        if count == 2:
            done_enter.set()

    n.start_soon(call_and_signal, ca.__aenter__)
    n.start_soon(call_and_signal, cb.__aenter__)
    await done_enter.wait()
    yield

    n.start_soon(ca.__aexit__)
    n.start_soon(cb.__aexit__)

The above doesn't work and triggers:

13:23:55,032 belt.sentry:185 88389   File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/trio/_core/_run.py", line 2084, in run
13:23:55,032 belt.sentry:185 88389     timeout = gen.send(next_send)
13:23:55,032 belt.sentry:185 88389               ^^^^^^^^^^^^^^^^^^^
13:23:55,032 belt.sentry:185 88389   File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/trio/_core/_run.py", line 2416, in unrolled_run
13:23:55,032 belt.sentry:185 88389     raise TrioInternalError("internal error in Trio - please file a bug!") from exc
13:23:55,032 belt.sentry:185 88389 trio.TrioInternalError: internal error in Trio - please file a bug!

(Sometimes it raises GeneratorExit instead.)

Rather than dig into that, here is a simpler version that just deals with a single context manager instead of two:

@asynccontextmanager
async def par_w(a):
    ca = a()
    await ca.__aenter__()
    yield
    await ca.__aexit__(None, None, None)

It raises the same:

13:25:08,219 belt.sentry:185 88488   File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/trio/_core/_run.py", line 2084, in run
13:25:08,219 belt.sentry:185 88488     timeout = gen.send(next_send)
13:25:08,219 belt.sentry:185 88488               ^^^^^^^^^^^^^^^^^^^
13:25:08,219 belt.sentry:185 88488   File "/opt/homebrew/Caskroom/miniforge/base/envs/py311/lib/python3.11/site-packages/trio/_core/_run.py", line 2416, in unrolled_run
13:25:08,219 belt.sentry:185 88488     raise TrioInternalError("internal error in Trio - please file a bug!") from exc
13:25:08,219 belt.sentry:185 88488 trio.TrioInternalError: internal error in Trio - please file a bug!

(Sometimes it also raises GeneratorExit instead.)

In the end, I thought of an entirely different approach—launch each context manager in its own task and synchronize with the original context manager to yield—but I was curious to understand why the above doesn't work. What trio contract is being violated?

[CoolCat467]

For line number debugging purposes, which exact version of Trio are you using? For example, on my machine, I can do this to get the version:

>>> import trio
>>> trio.__version__
'0.26.0'

[A5rocks]

Nevermind I misread this issue.

[yang]

0.22.2

[CoolCat467]

I've got this based off your minimal example:

from __future__ import annotations

import trio

from contextlib import asynccontextmanager, AbstractAsyncContextManager
from collections.abc import Callable, AsyncGenerator
from typing import TypeVar


T = TypeVar("T")


@asynccontextmanager
async def par_w(a: Callable[[], AbstractAsyncContextManager[T]]) -> AsyncGenerator[T, None]:
    ca = a()
    value = await ca.__aenter__()
    yield value
    await ca.__aexit__(None, None, None)


async def run() -> None:
    context_manager = trio.open_nursery
    async def do_three() -> None:
        print("3 start")
        await trio.sleep(1)
        print("3 end")
    async with par_w(context_manager) as nursery:
        print("1")
        await trio.sleep(1)
        print("2")
        nursery.start_soon(do_three)
        print("end")
    print("4")


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

and I don't see any errors happening, so I'm guessing it's an issue with whatever async context manager you are passing in as a to your par_w function. If you could share what that is that would be great.

[TeamSpen210]

The issue here is similar to the async generator issue. Because you're calling with statement methods directly, you have to be incredibly careful to preserve the nesting order. Otherwise, Trio's internal cancel scope stack can get messed up, breaking everything.

It's probably best to avoid calling them directly, instead use ExitStack, or put things in different tasks for concurrency.

[A5rocks]

Otherwise, Trio's internal cancel scope stack can get messed up, breaking everything.

That should raise a RuntimeError, not a TrioInternalError (though I guess I doubt Trio handles every possible case... but that seems unlikely)... At least that's what this exits with:

import trio
from contextlib import asynccontextmanager

@asynccontextmanager
async def par_w(a):
    ca = a()
    await ca.__aenter__()
    yield
    await ca.__aexit__(None, None, None)

async def printer():
    while True:
        print("printer printing")
        await trio.sleep(1)

@asynccontextmanager
async def acm():
    async with trio.open_nursery() as nursery:
        nursery.start_soon(printer)
        yield

async def main():
    with trio.move_on_after(5):
        async with par_w(acm):
            await trio.sleep(10)

trio.run(main)

(the bug here is that ca.__aexit__(None, None, None) is (obviously) wrong in case of exceptions such as cancellation. If you replace it as in https://peps.python.org/pep-0492/#new-syntax then it works.)

---

However that doesn't seem to be what this issue is about.

@yang could you provide any code that can reproduce the TrioInternalError (on latest Trio versions, preferably.)? I'll close this for now because without that this issue is not actionable.