Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Open
yang opened this issue Jul 14, 2024 · 5 comments

Comments

@yang
Copy link

yang commented Jul 14, 2024

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?

@yang yang changed the title Right way to implement parallel context manager combinator Why can't you await __aenter__? (Understanding trio's contracts around async context managers) Jul 14, 2024
@CoolCat467
Copy link
Member

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

This comment has been minimized.

@yang
Copy link
Author

yang commented Jul 15, 2024

0.22.2

@CoolCat467
Copy link
Member

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
Copy link
Contributor

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

4 participants