Leaking a timeout to the outer scope

Suppose that we want to iterate over an async iterator, but wait for at most max_time seconds for each element. We might naturally encapsulate the logic for doing so in an async generator, so that the call site can continue to use a straightforward async for loop:

async def iter_with_timeout(ait, max_time):
    try:
        while True:
            with timeout(max_time):
                yield await anext(ait)
    except StopAsyncIteration:
        return

async def fn():
    async for elem in iter_with_timeout(ait, max_time=1.0):
        await do_something_with(elem)

Unfortunately, there’s a bug in this version: the timeout might expire after the generator yields but before it is resumed! In this case, we’ll see a CancelledError raised in the outer task, where it cannot be caught by the with timeout(max_time): statement.

The fix is fairly simple: get the next element inside the timeout context, and then yield outside that context.

async def correct_iter_with_timeout(ait, max_time):
    try:
        while True:
            with timeout(max_time):
                tmp = await anext(ait)
            yield tmp
    except StopAsyncIteration:
        return

Leaking background tasks (breaks cancellation and exception handling)

Timeouts are not the only interface which wrap a cancel scope - and if you need some background worker tasks, you can’t simply close the TaskGroup before yielding.

As an example, let’s look at a fan-in generator, which we’ll use to merge the feeds from several “sensors”. We’ll also set up our mock sensors with a small buffer, so that we’ll raise an error in the background task while control flow is outside the combined_iterators generator.

import asyncio, itertools

async def mock_sensor(name):
    for n in itertools.count():
        await asyncio.sleep(0.1)
        if n == 1 and name == "b":  # 'presence detection'
            yield "PRESENT"
        elif n == 3 and name == "a":  # inject a simple bug
            print("oops, raising RuntimeError")
            raise RuntimeError
        else:
            yield f"{name}-{n}"  # non-presence sensor data

async def move_elements_to_queue(ait, queue):
    async for obj in ait:
        await queue.put(obj)

async def combined_iterators(*aits):
    """Combine async iterators by starting N tasks, each of
    which move elements from one iterable to a shared queue."""
    q = asyncio.Queue(maxsize=2)
    async with asyncio.TaskGroup() as tg:
        for ait in aits:
            tg.create_task(move_elements_to_queue(ait, q))
        while True:
            yield await q.get()

async def turn_on_lights_when_someone_gets_home():
    combined = combined_iterators(mock_sensor("a"), mock_sensor("b"))
    async for event in combined:
        print(event)
        if event == "PRESENT":
            break
    print("main task sleeping for a bit")
    await asyncio.sleep(1)  # do some other operation

asyncio.run(turn_on_lights_when_someone_gets_home())

When we run this code, we see the expected sequence of observations, then a ‘detection’, and then while the main task is sleeping we trigger that RuntimeError in the background. But… we don’t actually observe the RuntimeError, not even as the __context__ of another exception!

>> python3.11 demo.py
a-0
b-0
a-1
PRESENT
main task sleeping for a bit
oops, raising RuntimeError

Traceback (most recent call last):
  File "demo.py", line 39, in <module>
    asyncio.run(turn_on_lights_when_someone_gets_home())
  ...
  File "demo.py", line 37, in turn_on_lights_when_someone_gets_home
    await asyncio.sleep(1)  # do some other operation
  File ".../python3.11/asyncio/tasks.py", line 649, in sleep
    return await future
asyncio.exceptions.CancelledError

Here, again, the problem is that we’ve yielded inside a cancel scope; this time the scope which a TaskGroup uses to cancel sibling tasks when one of the child tasks raises an exception. However, the CancelledError which was intended for the sibling task was instead injected into the outer task, and so we never got a chance to create and raise an ExceptionGroup(..., [RuntimeError()]).

To fix this, we need to turn our async generator into an async context manager, which yields an async iterable - in this case a generator wrapping the queue; in future perhaps the queue itself:

async def queue_as_aiterable(queue):
    # async generators that don't `yield` inside a cancel scope are fine!
    while True:
        try:
            yield await queue.get()
        except asyncio.QueueShutDown:
            return

@asynccontextmanager  #  yield-in-cancel-scope is OK in a context manager
async def combined_iterators(*aits):
    q = asyncio.Queue(maxsize=2)
    async with asyncio.TaskGroup() as tg:
        for ait in aits:
            tg.create_task(move_elements_to_queue(ait, q))
        yield queue_as_aiterable(q)

async def turn_on_lights_when_someone_gets_home():
    ...
    async with combined_iterators(...) as ait:
        async for event in ait:
            ...

In a user-defined context manager

Yielding inside a cancel scope can be safe, if and only if you’re using the generator to implement a context manager [3] - in this case any propagating exceptions will be redirected to the expected task.

We’ve also implemented the ASYNC101 linter rule in flake8-async, which warns against yielding in known cancel scopes. Could user education be sufficient to avoid these problems? Unfortunately not: user-defined context managers can also wrap a cancel scope, and it’s infeasible to recognize or lint for all such cases.

This regularly arises in practice, because ‘run some background tasks for the duration of this context’ is a very common pattern in structured concurrency. We saw that in combined_iterators() above; and have seen this bug in multiple implementations of the websocket protocol:

async def get_messages(websocket_url):
    # The websocket protocol requires background tasks to manage the socket heartbeat
    async with open_websocket(websocket_url) as ws:  # contains a TaskGroup!
        while True:
            yield await ws.get_message()

async with open_websocket(websocket_url) as ws:
    async for message in get_messages(ws):
        ...

Implementation - tracking frames

The new sys.prevent_yields context manager will require interpreter support. For each frame, we track the entries and exits of this context manager.

We’re not particularly attached to the exact representation; we’ll discuss it as a stack (which would support clear error messages), but more compact representations such as pair-of-integers would also work.

• When entering a newly-created or resumed frame, initialize empty stacks of entries and exits.
• When returning from a frame, merge these stacks into that of the parent frame.
• When yielding:
  • if entries != [] and not frame.allow_yield_flag, raise a RuntimeError instead of yielding (the new behavior this PEP proposes)
  • otherwise, merge stacks into the parent frame as for a return.

Because this is about yielding frames within a task, not switching between tasks, syntactic yield and yield from should be affected, but await expressions should not.

We can reduce the overhead by storing this metadata in a single stack per thread for all stack frames which are not generators.

Worked examples

No-yield example

In this example, we see multiple rounds of the stack merging as we unwind from sys.prevent_yields, through the user-defined ContextManager, back to the original Frame. For brevity, the reason for preventing yields is not shown; it is part of the “1 enter” state.

With no yield we don’t raise any errors, and because the number of enters and exits balance the frame returns as usual with no further tracking.

Attempts-to-yield example

In this example, the Frame attempts to yield while inside the sys.prevent_yields context. This is detected by the interpreter, which raises a RuntimeError instead of suspending the frame.

Allowed-to-yield example

In this example, a decorator has marked the Frame as allowing yields. This could be @contextlib.contextmanager or a related decorator.

When the Frame is allowed to yield, the entry/exit stack is merged into the parent frame’s stack before suspending. When the Frame resumes, its stack is empty. Finally, when the Frame exits, the exit is merged into the parent frame’s stack, rebalancing it.

This ensures that the parent frame correctly inherits any remaining sys.prevent_yields state, while allowing the Frame to safely suspend and resume.

Behavior if sys.prevent_yields is misused

While unwise, it’s possible to call sys.prevent_yields.__enter__ and .__exit__ in an order that does not correspond to any valid nesting, or get an invalid frame state in some other way.

There are two ways sys.prevent_yields.__exit__ could detect an invalid state. First, if yields are not prevented, we can simply raise an exception without changing the state. Second, if an unexpected entry is at the top of the stack, we suggest popping that entry and raising an exception – this ensures that out-of-order calls will still clear the stack, while still making it clear that something is wrong.

(and if we choose e.g. an integer- rather than stack-based representation, such states may not be distinguishable from correct nesting at all, in which case the question will not arise)

How widespread is this bug?

We don’t have solid numbers here, but believe that many projects are affected in the wild. Since hitting a moderate and a critical bug attributed to suspending a cancel scope in the same week at work, we’ve used static analysis with some success. Three people Zac spoke to at PyCon recognized the symptoms and concluded that they had likely been affected.

TODO: run the ASYNC101 lint rule across ecosystem projects, e.g. the aio-libs packages, and get some sense of frequency in widely-used PyPI packages? This would help inform the break/deprecation pathways for stdlib code.

PEP 533, deterministic cleanup for iterators

PEP 533 proposes adding __[a]iterclose__ to the iterator protocol, essentially wrapping a with [a]closing(ait) around each (async) for loop. While this would be useful for ensuring timely and deterministic cleanup of resources held by iterators, the problem it aims to solve, it does not fully address the issues that motivate this PEP.

Even with PEP 533, misfired cancellations would still be delivered to the wrong task and could wreak havoc before the iterator is closed. Moreover, it does not address the fundamental structured concurrency problem with TaskGroup, where suspending a frame that owns a TaskGroup is incompatible with the model of child tasks being fully encapsulated within their parent frame.

Deprecate async generators entirely

At the 2024 language summit, several attendees suggested instead deprecating async generators in toto. Unfortunately, while the common-in-practice cases all use async generators, Trio code can trigger the same problem with standard generators:

# We use Trio for this example, because while `asyncio.timeout()` is async,
# Trio's CancelScope type and timeout context managers are synchronous.
import trio

def abandon_each_iteration_after(max_seconds):
    # This is of course broken, but I can imagine someone trying it...
    while True:
        with trio.move_on_after(max_seconds):
            yield

@trio.run
async def main():
    for _ in abandon_each_iteration_after(max_seconds=1):
        await trio.sleep(3)

If it wasn’t for the bug in question, this code would look pretty idiomatic - but after about a second, instead of moving on to the next iteration it raises:

Traceback (most recent call last):
  File "demo.py", line 10, in <module>
    async def main():
  File "trio/_core/_run.py", line 2297, in run
    raise runner.main_task_outcome.error
  File "demo.py", line 12, in main
    await trio.sleep(3)
  File "trio/_timeouts.py", line 87, in sleep
    await sleep_until(trio.current_time() + seconds)
  ...
  File "trio/_core/_run.py", line 1450, in raise_cancel
    raise Cancelled._create()
trio.Cancelled: Cancelled

Furthermore, there are some non-cancel-scope synchronous context managers which exhibit related problems, such as the abovementioned decimal.localcontext. While fixing the example below is not a goal of this PEP, it demonstrates that yield-within-with problems are not exclusive to async generators:

import decimal

def why_would_you_do_this():
    with decimal.localcontext(decimal.Context(prec=1)):
        yield

one = decimal.Decimal(1)
print(one / 3)  # 0.3333333333333333333333333333
next(gen := why_would_you_do_this())
print(one / 3)  # 0.3

While I’ve had good experiences in async Python without async generators [5], I’d prefer to fix the problem than remove them from the language.

Can’t we just deliver exceptions to the right place?

If we implemented PEP 568 (Generator-sensitivity for Context Variables; see also PEP 550), it would be possible to handle exceptions from timeouts: the event loop could avoid firing a CancelledError until the generator frame which contains the context manager is on the stack - either when the generator is resumed, or when it is finalized.

This can take arbitrarily long; even if we implemented PEP 533 to ensure timely cleanup on exiting (async) for-loops it’s still possible to drive a generator manually with next/send.

However, this doesn’t address the other problem with TaskGroup. The model for generators is that you put a stack frame in suspended animation and can then treat it as an inert value which can be stored, moved around, and maybe discarded or revived in some arbitrary place. The model for structured concurrency is that your stack becomes a tree, with child tasks encapsulated within some parent frame. They’re extending the basic structured programming model in different, and unfortunately incompatible, directions.

Suppose for example that suspending a frame containing an open TaskGroup also suspended all child tasks. This would preserve the ‘downward’ structured concurrency, in that children remain encapsulated - albeit at the cost of deadlocking both of our motivating examples, and much real-world code. However, it would still be possible to resume the generator in a different task, violating the ‘upwards’ invariant of structured concurrency.

We don’t think it’s worth adding this much machinery to handle cancel scopes, while still leaving task groups broken.

Alternative implementation - inspecting bytecode

Jelle Zijlstra has sketched an alternative, where sys.prevent_yields inspects the bytecode of callers until satisfied that there is no yield between the calling instruction pointer and the next context exit. We expect that support for syntatically-nested context managers could be added fairly easily.

However, it’s not yet clear how this would work when user-defined context managers wrap sys.prevent_yields. Worse, this approach ignores explicit calls to __enter__() and __exit__(), meaning that the context management protocol would vary depending on whether the with statement was used.

The ‘only pay if you use it’ performance cost is very attractive. However, inspecting frame objects is prohibitively expensive for core control-flow constructs, and causes whole-program slowdowns via de-optimization. On the other hand, adding interpreter support for better performance leads back to the same pay-regardless semantics as our preferred solution above.