Skip to content

Commit

Permalink
Merge branch 'dont-use-f_locals-for-agen-finalization' of github.com:…
Browse files Browse the repository at this point in the history
…graingert/trio into avoid-refcycles-in-run-exc
  • Loading branch information
graingert committed Oct 26, 2024
2 parents 1e10e62 + c4010fc commit 6bf1421
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 18 deletions.
5 changes: 5 additions & 0 deletions newsfragments/3112.bugfix.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
Rework foreign async generator finalization to track async generator
ids rather than mutating ``ag_frame.f_locals``. This fixes an issue
with the previous implementation: locals' lifetimes will no longer be
extended by materialization in the ``ag_frame.f_locals`` dictionary that
the previous finalization dispatcher logic needed to access to do its work.
18 changes: 11 additions & 7 deletions src/trio/_core/_asyncgens.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,11 @@ class AsyncGenerators:
# regular set so we don't have to deal with GC firing at
# unexpected times.
alive: _WEAK_ASYNC_GEN_SET | _ASYNC_GEN_SET = attrs.Factory(_WEAK_ASYNC_GEN_SET)
# The ids of foreign async generators are added to this set when first
# iterated. Usually it is not safe to refer to ids like this, but because
# we're using a finalizer we can ensure ids in this set do not outlive
# their async generator.
foreign: set[int] = attrs.Factory(set)

# This collects async generators that get garbage collected during
# the one-tick window between the system nursery closing and the
Expand All @@ -51,10 +56,7 @@ def firstiter(agen: AsyncGeneratorType[object, NoReturn]) -> None:
# An async generator first iterated outside of a Trio
# task doesn't belong to Trio. Probably we're in guest
# mode and the async generator belongs to our host.
# The locals dictionary is the only good place to
# remember this fact, at least until
# https://bugs.python.org/issue40916 is implemented.
agen.ag_frame.f_locals["@trio_foreign_asyncgen"] = True
self.foreign.add(id(agen))
if self.prev_hooks.firstiter is not None:
self.prev_hooks.firstiter(agen)

Expand All @@ -77,12 +79,14 @@ def finalize_in_trio_context(
self.trailing_needs_finalize.add(agen)

def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None:
agen_name = name_asyncgen(agen)
try:
is_ours = not agen.ag_frame.f_locals.get("@trio_foreign_asyncgen")
except AttributeError: # pragma: no cover
self.foreign.remove(id(agen))
except KeyError:
is_ours = True
else:
is_ours = False

agen_name = name_asyncgen(agen)
if is_ours:
runner.entry_queue.run_sync_soon(
finalize_in_trio_context,
Expand Down
64 changes: 53 additions & 11 deletions src/trio/_core/_tests/test_guest_mode.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

import asyncio
import contextlib
import contextvars
import queue
import signal
import socket
Expand All @@ -11,6 +10,7 @@
import time
import traceback
import warnings
import weakref
from collections.abc import AsyncGenerator, Awaitable, Callable
from functools import partial
from math import inf
Expand All @@ -22,6 +22,7 @@
)

import pytest
import sniffio
from outcome import Outcome

import trio
Expand Down Expand Up @@ -221,7 +222,8 @@ async def trio_main(in_host: InHost) -> str:


def test_guest_mode_sniffio_integration() -> None:
from sniffio import current_async_library, thread_local as sniffio_library
current_async_library = sniffio.current_async_library
sniffio_library = sniffio.thread_local

async def trio_main(in_host: InHost) -> str:
async def synchronize() -> None:
Expand Down Expand Up @@ -439,9 +441,9 @@ def aiotrio_run(
loop = asyncio.new_event_loop()

async def aio_main() -> T:
trio_done_fut = loop.create_future()
trio_done_fut: asyncio.Future[Outcome[T]] = loop.create_future()

def trio_done_callback(main_outcome: Outcome[object]) -> None:
def trio_done_callback(main_outcome: Outcome[T]) -> None:
print(f"trio_fn finished: {main_outcome!r}")
trio_done_fut.set_result(main_outcome)

Expand All @@ -455,9 +457,11 @@ def trio_done_callback(main_outcome: Outcome[object]) -> None:
**start_guest_run_kwargs,
)

return (await trio_done_fut).unwrap() # type: ignore[no-any-return]
return (await trio_done_fut).unwrap()

try:
# can't use asyncio.run because that fails on Windows (3.8, x64, with
# Komodia LSP)
return loop.run_until_complete(aio_main())
finally:
loop.close()
Expand Down Expand Up @@ -628,8 +632,6 @@ async def trio_main(in_host: InHost) -> None:

@restore_unraisablehook()
def test_guest_mode_asyncgens() -> None:
import sniffio

record = set()

async def agen(label: str) -> AsyncGenerator[int, None]:
Expand All @@ -656,9 +658,49 @@ async def trio_main() -> None:

gc_collect_harder()

# Ensure we don't pollute the thread-level context if run under
# an asyncio without contextvars support (3.6)
context = contextvars.copy_context()
context.run(aiotrio_run, trio_main, host_uses_signal_set_wakeup_fd=True)
aiotrio_run(trio_main, host_uses_signal_set_wakeup_fd=True)

assert record == {("asyncio", "asyncio"), ("trio", "trio")}


@restore_unraisablehook()
def test_guest_mode_asyncgens_garbage_collection() -> None:
record: set[tuple[str, str, bool]] = set()

async def agen(label: str) -> AsyncGenerator[int, None]:
class A:
pass

a = A()
a_wr = weakref.ref(a)
assert sniffio.current_async_library() == label
try:
yield 1
finally:
library = sniffio.current_async_library()
with contextlib.suppress(trio.Cancelled):
await sys.modules[library].sleep(0)

del a
if sys.implementation.name == "pypy":
gc_collect_harder()

record.add((label, library, a_wr() is None))

async def iterate_in_aio() -> None:
await agen("asyncio").asend(None)

async def trio_main() -> None:
task = asyncio.ensure_future(iterate_in_aio())
done_evt = trio.Event()
task.add_done_callback(lambda _: done_evt.set())
with trio.fail_after(1):
await done_evt.wait()

await agen("trio").asend(None)

gc_collect_harder()

aiotrio_run(trio_main, host_uses_signal_set_wakeup_fd=True)

assert record == {("asyncio", "asyncio", True), ("trio", "trio", True)}

0 comments on commit 6bf1421

Please sign in to comment.