diff --git a/newsfragments/3554.bugfix.rst b/newsfragments/3554.bugfix.rst new file mode 100644 index 0000000000..edd5efa642 --- /dev/null +++ b/newsfragments/3554.bugfix.rst @@ -0,0 +1,4 @@ +When an ``async with`` on an `@as_safe_channel ` generator is cancelled +while the generator's own code is running, the generator's `~trio.Cancelled` now propagates, as +it would from a plain ``async for``. Previously the choice was random, so an exception the +generator raised before its cleanup was cancelled could be missing from the traceback. diff --git a/src/trio/_channel.py b/src/trio/_channel.py index c785434e7a..da3c3afbda 100644 --- a/src/trio/_channel.py +++ b/src/trio/_channel.py @@ -22,6 +22,7 @@ MultipleExceptionError, NoPublicConstructor, final, + raise_saving_context, raise_single_exception_from_group, ) @@ -559,6 +560,7 @@ async def context_manager( *args: P.args, **kwargs: P.kwargs ) -> AsyncGenerator[trio._channel.RecvChanWrapper[T], None]: send_chan, recv_chan = trio.open_memory_channel[T](0) + agen_cancelled: list[trio.Cancelled] = [] try: async with trio.open_nursery(strict_exception_groups=True) as nursery: agen = fn(*args, **kwargs) @@ -567,7 +569,11 @@ async def context_manager( # If this errors we don't close `recv_chan`, but the caller # never gets access to it, so that's not a problem. await nursery.start( - _move_elems_to_channel, agen, send_chan, send_semaphore + _move_elems_to_channel, + agen, + send_chan, + send_semaphore, + agen_cancelled, ) # `async with recv_chan` could eat exceptions, so use sync cm with RecvChanWrapper(recv_chan, send_semaphore) as wrapped_recv_chan: @@ -578,6 +584,13 @@ async def context_manager( "exited trio.as_safe_channel context manager" ) except BaseExceptionGroup as eg: + # If the generator's own code was cancelled and the group holds nothing but + # Cancelled, raise what a plain `async for` would: the generator's Cancelled, + # with its context. (A group raised by the body is not a Cancelled.) + if agen_cancelled and all( + isinstance(e, trio.Cancelled) for e in eg.exceptions + ): + raise_saving_context(agen_cancelled[0]) try: raise_single_exception_from_group(eg) except MultipleExceptionError: @@ -590,11 +603,16 @@ async def context_manager( ) raise eg from None + finally: + # Break the reference cycle: the Cancelled's traceback holds this frame, + # which holds the list. + agen_cancelled.clear() async def _move_elems_to_channel( agen: AsyncGenerator[T, None], send_chan: trio.MemorySendChannel[T], send_semaphore: trio.Semaphore, + agen_cancelled: list[trio.Cancelled], task_status: trio.TaskStatus, ) -> None: # `async with send_chan` will eat exceptions, @@ -615,6 +633,9 @@ async def _move_elems_to_channel( value = await agen.__anext__() except StopAsyncIteration: return + except trio.Cancelled as cancelled: + agen_cancelled.append(cancelled) + raise # Send the value to the channel try: await send_chan.send(value) diff --git a/src/trio/_tests/test_channel.py b/src/trio/_tests/test_channel.py index 831f8eb483..2a609393fe 100644 --- a/src/trio/_tests/test_channel.py +++ b/src/trio/_tests/test_channel.py @@ -638,6 +638,32 @@ async def agen(events: list[str]) -> AsyncGenerator[None]: assert events == ["body cancel", "agen cancel"] +async def test_as_safe_channel_cancelled_in_cleanup_keeps_context( + autojump_clock: trio.abc.Clock, +) -> None: + error = ValueError("agen") + + @as_safe_channel + async def agen() -> AsyncGenerator[None]: + try: + yield + raise error + finally: + # The timeout cancels the body and this task. The shield makes this task's + # Cancelled, which carries `error`, reach the nursery after the body's; + # unshielded, the order is random. + with trio.CancelScope(shield=True): + await trio.sleep(2) + await trio.lowlevel.checkpoint() + + with trio.move_on_after(1): + with pytest.raises(trio.Cancelled) as excinfo: + async with agen() as recv_chan: + async for _ in recv_chan: # pragma: no branch + pass + assert excinfo.value.__context__ is error + + async def test_as_safe_channel_genexit_exception_group() -> None: @as_safe_channel async def agen() -> AsyncGenerator[None]: diff --git a/src/trio/_util.py b/src/trio/_util.py index bae6608d88..55cd03a119 100644 --- a/src/trio/_util.py +++ b/src/trio/_util.py @@ -358,7 +358,7 @@ def raise_single_exception_from_group( the cause/context of the group(s). If multiple non-cancelled exceptions are encountered, it raises - :exc:`AssertionError`. + :exc:`MultipleExceptionError`. """ # immediately bail out if there's any KI or SystemExit for e in eg.exceptions: