Leaking a timeout to the outer scope

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

asyncdefiter_with_timeout(ait,max_time):try:whileTrue:withtimeout(max_time):yieldawaitanext(ait)exceptStopAsyncIteration:returnasyncdeffn():asyncforeleminiter_with_timeout(ait,max_time=1.0):awaitdo_something_with(elem)

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

The fix is fairly simple: get the next element inside the timeout context, andthen yieldoutside that context.

asyncdefcorrect_iter_with_timeout(ait,max_time):try:whileTrue:withtimeout(max_time):tmp=awaitanext(ait)yieldtmpexceptStopAsyncIteration:return

Leaking background tasks (breaks cancellation and exception handling)

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

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

importasyncio,itertoolsasyncdefmock_sensor(name):forninitertools.count():awaitasyncio.sleep(0.1)ifn==1andname=="b":# 'presence detection'yield"PRESENT"elifn==3andname=="a":# inject a simple bugprint("oops, raising RuntimeError")raiseRuntimeErrorelse:yieldf"{name}-{n}"# non-presence sensor dataasyncdefmove_elements_to_queue(ait,queue):asyncforobjinait:awaitqueue.put(obj)asyncdefcombined_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)asyncwithasyncio.TaskGroup()astg:foraitinaits:tg.create_task(move_elements_to_queue(ait,q))whileTrue:yieldawaitq.get()asyncdefturn_on_lights_when_someone_gets_home():combined=combined_iterators(mock_sensor("a"),mock_sensor("b"))asyncforeventincombined:print(event)ifevent=="PRESENT":breakprint("main task sleeping for a bit")awaitasyncio.sleep(1)# do some other operationasyncio.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 thatRuntimeError in the background. But… we don’t actually observe theRuntimeError, not even as the__context__ of another exception!

>> python3.11 demo.pya-0b-0a-1PRESENTmain task sleeping for a bitoops, raising RuntimeErrorTraceback (most recent call last):  File"demo.py", line39, in<module>asyncio.run(turn_on_lights_when_someone_gets_home())...  File"demo.py", line37, inturn_on_lights_when_someone_gets_homeawaitasyncio.sleep(1)# do some other operation  File".../python3.11/asyncio/tasks.py", line649, insleepreturnawaitfutureasyncio.exceptions.CancelledError

Here, again, the problem is that we’veyielded inside a cancel scope;this time the scope which aTaskGroup uses to cancel sibling tasks when oneof the child tasks raises an exception. However, theCancelledError whichwas intended for the sibling task was instead injected into theouter task,and so we never got a chance to create and raise anExceptionGroup(...,[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; infutureperhaps the queue itself:

asyncdefqueue_as_aiterable(queue):# async generators that don't `yield` inside a cancel scope are fine!whileTrue:try:yieldawaitqueue.get()exceptasyncio.QueueShutDown:return@asynccontextmanager#  yield-in-cancel-scope is OK in a context managerasyncdefcombined_iterators(*aits):q=asyncio.Queue(maxsize=2)asyncwithasyncio.TaskGroup()astg:foraitinaits:tg.create_task(move_elements_to_queue(ait,q))yieldqueue_as_aiterable(q)asyncdefturn_on_lights_when_someone_gets_home():...asyncwithcombined_iterators(...)asait:asyncforeventinait:...

In a user-defined context manager

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

We’ve also implemented theASYNC101 linter rule inflake8-async, which warns against yielding inknown cancel scopes. Could user education be sufficient to avoid theseproblems? Unfortunately not: user-defined context managers can also wrap acancel scope, and it’s infeasible to recognize or lint for all such cases.

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

asyncdefget_messages(websocket_url):# The websocket protocol requires background tasks to manage the socket heartbeatasyncwithopen_websocket(websocket_url)asws:# contains a TaskGroup!whileTrue:yieldawaitws.get_message()asyncwithopen_websocket(websocket_url)asws:asyncformessageinget_messages(ws):...

Implementation - tracking frames

The newsys.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 asa stack (which would support clear error messages), but more compactrepresentations such as pair-of-integers would also work.

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

Because this is about yielding frameswithin a task, not switching betweentasks, syntacticyield andyieldfrom should be affected, butawaitexpressions should not.

We can reduce the overhead by storing this metadata in a single stack per threadfor 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 fromsys.prevent_yields, through the user-defined ContextManager, back to theoriginal Frame. For brevity, the reason for preventing yields is not shown;it is part of the “1 enter” state.

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

Attempts-to-yield example

In this example, the Frame attempts toyield while inside thesys.prevent_yields context. This is detected by the interpreter,which raises aRuntimeError instead of suspending the frame.

Allowed-to-yield example

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

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

This ensures that the parent frame correctly inherits any remainingsys.prevent_yields state, while allowing the Frame to safely suspendand resume.

Behavior ifsys.prevent_yields is misused

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

There are two wayssys.prevent_yields.__exit__ could detect an invalid state.First, if yields are not prevented, we can simply raise an exception withoutchanging 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 thatout-of-order calls will still clear the stack, while still making it clear thatsomething is wrong.

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

How widespread is this bug?

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

TODO: run the ASYNC101 lint rule across ecosystem projects, e.g. the aio-libspackages, 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 awith[a]closing(ait) around each (async) for loop.While this would be useful for ensuring timely and deterministic cleanup ofresources held by iterators, the problem it aims to solve, it does not fullyaddress the issues that motivate this PEP.

Even with PEP 533, misfired cancellations would still be delivered to the wrongtask and could wreak havoc before the iterator is closed. Moreover, it does notaddress the fundamental structured concurrency problem withTaskGroup, wheresuspending a frame that owns a TaskGroup is incompatible with the model of childtasks being fully encapsulated within their parent frame.

Deprecate async generators entirely

At the 2024 language summit, several attendees suggested instead deprecating asyncgeneratorsin toto. Unfortunately, while the common-in-practice cases all useasync 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.importtriodefabandon_each_iteration_after(max_seconds):# This is of course broken, but I can imagine someone trying it...whileTrue:withtrio.move_on_after(max_seconds):yield@trio.runasyncdefmain():for_inabandon_each_iteration_after(max_seconds=1):awaittrio.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", line10, in<module>asyncdefmain():  File"trio/_core/_run.py", line2297, inrunraiserunner.main_task_outcome.error  File"demo.py", line12, inmainawaittrio.sleep(3)  File"trio/_timeouts.py", line87, insleepawaitsleep_until(trio.current_time()+seconds)...  File"trio/_core/_run.py", line1450, inraise_cancelraiseCancelled._create()trio.Cancelled:Cancelled

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

importdecimaldefwhy_would_you_do_this():withdecimal.localcontext(decimal.Context(prec=1)):yieldone=decimal.Decimal(1)print(one/3)# 0.3333333333333333333333333333next(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 thelanguage.

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

If we implementedPEP 568 (Generator-sensitivity for Context Variables; seealsoPEP 550), it would be possible to handle exceptions from timeouts: theevent loop could avoid firing aCancelledError until the generator framewhich contains the context manager is on the stack - either when the generatoris resumed, or when it is finalized.

This can take arbitrarily long; even if we implementedPEP 533 to ensuretimely cleanup on exiting (async) for-loops it’s still possible to drive agenerator manually with next/send.

However, this doesn’t address the other problem withTaskGroup. The modelfor generators is that you put a stack frame in suspended animation and can thentreat it as an inert value which can be stored, moved around, and maybediscarded or revived in some arbitrary place. The model for structuredconcurrency is that your stack becomes a tree, with child tasks encapsulatedwithin some parent frame. They’re extending the basic structured programmingmodel in different, and unfortunately incompatible, directions.

Suppose for example that suspending a frame containing an openTaskGroupalso suspended all child tasks. This would preserve the ‘downward’ structuredconcurrency, in that children remain encapsulated - albeit at the cost ofdeadlocking both of our motivating examples, and much real-world code.However, it would still be possible to resume the generator in a differenttask, 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 hassketched an alternative, wheresys.prevent_yieldsinspects the bytecode of callers until satisfied that there is no yield betweenthe calling instruction pointer and the next context exit. We expect thatsupport for syntatically-nested context managers could be added fairly easily.

However, it’s not yet clear how this would work when user-defined contextmanagers wrapsys.prevent_yields. Worse, this approach ignores explicitcalls to__enter__() and__exit__(), meaning that the context managementprotocol would vary depending on whether thewith 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-flowconstructs, and causes whole-program slowdowns via de-optimization.On the other hand, adding interpreter support for better performance leadsback to the same pay-regardless semantics as our preferred solution above.