async def owned_stream(source: AsyncGenerator[_T, None]) -> AsyncGenerator[_T, None]:
"""Relay values and exceptions; join the source owner on every exit."""
commands: asyncio.Queue[tuple[str, Any]] = asyncio.Queue()
replies: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(maxsize=1)
phase = "waiting"
async def drive() -> None:
nonlocal phase
try:
while True:
phase = "waiting"
kind, payload = await commands.get()
phase = "advancing"
if kind == "close":
return
try:
value = (
await source.athrow(payload)
if kind == "throw"
else await anext(source)
)
except StopAsyncIteration:
phase = "closing"
await replies.put(("done", None))
return
except BaseException as exc:
phase = "closing"
await replies.put(("error", exc))
return
else:
await replies.put(("item", value))
finally:
phase = "closing"
await source.aclose()
owner = asyncio.create_task(drive(), name="symfonic-stream-owner")
terminal = False
close_requested = False
cancelled_exit = False
cancel_relayed = False
async def reply() -> tuple[str, Any]:
receive = asyncio.create_task(replies.get())
try:
done, _ = await asyncio.wait(
{owner, receive}, return_when=asyncio.FIRST_COMPLETED
)
if receive in done or not replies.empty():
return await receive
await owner
return ("done", None)
except asyncio.CancelledError:
if receive.done() and not receive.cancelled():
replies.put_nowait(receive.result())
raise
finally:
if not receive.done():
receive.cancel()
with contextlib.suppress(asyncio.CancelledError):
await receive
async def join() -> None:
# Once cancellation reaches the source, a source that suppresses it
# may yield again. Queue a close request and drain that response so it
# cannot block the owner before teardown.
while not owner.done():
receive = asyncio.create_task(replies.get())
done, _ = await asyncio.wait(
{owner, receive}, return_when=asyncio.FIRST_COMPLETED
)
if receive not in done:
receive.cancel()
with contextlib.suppress(asyncio.CancelledError):
await receive
break
kind, value = receive.result()
if kind == "error" and not isinstance(value, asyncio.CancelledError):
await owner
if cancelled_exit and not cancel_relayed:
raise asyncio.CancelledError from value
raise value
with contextlib.suppress(asyncio.CancelledError):
await owner
while not replies.empty():
kind, value = replies.get_nowait()
if kind == "error" and not isinstance(value, asyncio.CancelledError):
if cancelled_exit and not cancel_relayed:
raise asyncio.CancelledError from value
raise value
try:
command: tuple[str, Any] = ("next", None)
while True:
await commands.put(command)
kind, value = await reply()
if kind == "item":
try:
yield value
except GeneratorExit:
raise
except BaseException as exc:
task = asyncio.current_task()
if isinstance(exc, asyncio.CancelledError) and task and task.cancelling():
raise
command = ("throw", exc)
else:
command = ("next", None)
elif kind == "error":
terminal = True
raise value
else:
terminal = True
return
except asyncio.CancelledError:
cancelled_exit = not terminal
if not terminal:
if phase == "advancing" or not commands.empty():
cancel_relayed = owner.cancel()
commands.put_nowait(("close", None))
close_requested = True
raise
finally:
if not terminal and not close_requested and not owner.done():
commands.put_nowait(("close", None))
cleanup = asyncio.create_task(join())
interrupted = False
while not cleanup.done():
try:
await asyncio.shield(cleanup)
except asyncio.CancelledError:
interrupted = True
except BaseException:
break
cleanup.result()
if interrupted:
raise asyncio.CancelledError