Skip to content

Commit

Permalink
ki protect asyncgen finalizers, because it's critical to run
Browse files Browse the repository at this point in the history
  • Loading branch information
graingert committed Oct 27, 2024
1 parent 0ebad56 commit 28957e0
Showing 1 changed file with 28 additions and 19 deletions.
47 changes: 28 additions & 19 deletions src/trio/_core/_asyncgens.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,32 @@
_ASYNC_GEN_SET = set


@_core.disable_ki_protection
def _finalize_without_ki_protection(
agen_name: str,
agen: AsyncGeneratorType[object, NoReturn],
) -> None:
# Host has no finalizer. Reimplement the default
# Python behavior with no hooks installed: throw in
# GeneratorExit, step once, raise RuntimeError if
# it doesn't exit.
closer = agen.aclose()
try:
# If the next thing is a yield, this will raise RuntimeError
# which we allow to propagate
closer.send(None)
except StopIteration:
pass
else:
# If the next thing is an await, we get here. Give a nicer
# error than the default "async generator ignored GeneratorExit"
raise RuntimeError(
f"Non-Trio async generator {agen_name!r} awaited something "
"during finalization; install a finalization hook to "
"support this, or wrap it in 'async with aclosing(...):'",
)


@attrs.define(eq=False)
class AsyncGenerators:
# Async generators are added to this set when first iterated. Any
Expand Down Expand Up @@ -78,6 +104,7 @@ def finalize_in_trio_context(
# have hit it.
self.trailing_needs_finalize.add(agen)

@_core.enable_ki_protection
def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None:
try:
self.foreign.remove(id(agen))
Expand Down Expand Up @@ -112,25 +139,7 @@ def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None:
if self.prev_hooks.finalizer is not None:
self.prev_hooks.finalizer(agen)
else:
# Host has no finalizer. Reimplement the default
# Python behavior with no hooks installed: throw in
# GeneratorExit, step once, raise RuntimeError if
# it doesn't exit.
closer = agen.aclose()
try:
# If the next thing is a yield, this will raise RuntimeError
# which we allow to propagate
closer.send(None)
except StopIteration:
pass
else:
# If the next thing is an await, we get here. Give a nicer
# error than the default "async generator ignored GeneratorExit"
raise RuntimeError(
f"Non-Trio async generator {agen_name!r} awaited something "
"during finalization; install a finalization hook to "
"support this, or wrap it in 'async with aclosing(...):'",
)
_finalize_without_ki_protection(agen_name, agen)

self.prev_hooks = sys.get_asyncgen_hooks()
sys.set_asyncgen_hooks(firstiter=firstiter, finalizer=finalizer) # type: ignore[arg-type] # Finalizer doesn't use AsyncGeneratorType
Expand Down

0 comments on commit 28957e0

Please sign in to comment.