Skip to content

Commit 0ec9f8a

Browse files
committed
use a strong set of ids for foriegn async gens
1 parent be11787 commit 0ec9f8a

File tree

1 file changed

+9
-6
lines changed

1 file changed

+9
-6
lines changed

src/trio/_core/_asyncgens.py

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,11 @@ class AsyncGenerators:
3636
# regular set so we don't have to deal with GC firing at
3737
# unexpected times.
3838
alive: _WEAK_ASYNC_GEN_SET | _ASYNC_GEN_SET = attrs.Factory(_WEAK_ASYNC_GEN_SET)
39+
# The ids of foreign async generators are added to this set when first
40+
# iterated. Usually it is not safe to refer to ids like this, but because
41+
# we're using a finalizer we can ensure ids in this set do not outlive
42+
# their async generator.
43+
foreign: set[int] = attrs.Factory(set)
3944

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

@@ -80,8 +82,9 @@ def finalize_in_trio_context(
8082
def finalizer(agen: AsyncGeneratorType[object, NoReturn]) -> None:
8183
agen_name = name_asyncgen(agen)
8284
try:
83-
is_ours = not agen.ag_frame.f_locals.get("@trio_foreign_asyncgen")
84-
except AttributeError: # pragma: no cover
85+
self.foreign.remove(id(agen))
86+
is_ours = False
87+
except KeyError:
8588
is_ours = True
8689

8790
if is_ours:

0 commit comments

Comments
 (0)