[mypyc] Clear coroutine env on coroutine completion

mypyc/codegen/emitclass.py

@@ -252,7 +252,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
     getseters_name = f"{name_prefix}_getseters"
     vtable_name = f"{name_prefix}_vtable"
     traverse_name = f"{name_prefix}_traverse"
-    clear_name = f"{name_prefix}_clear"
+    clear_name = emitter.native_function_name(cl.clear)
     dealloc_name = f"{name_prefix}_dealloc"
     methods_name = f"{name_prefix}_methods"
     vtable_setup_name = f"{name_prefix}_trait_vtable_setup"
@@ -280,7 +280,7 @@ def generate_class(cl: ClassIR, module: str, emitter: Emitter) -> None:
         fields["tp_dealloc"] = f"(destructor){name_prefix}_dealloc"
         if not cl.is_acyclic:
             fields["tp_traverse"] = f"(traverseproc){name_prefix}_traverse"
-            fields["tp_clear"] = f"(inquiry){name_prefix}_clear"
+            fields["tp_clear"] = f"(inquiry){clear_name}"
     # Populate .tp_finalize and generate a finalize method only if __del__ is defined for this class.
     del_method = next((e.method for e in cl.vtable_entries if e.name == "__del__"), None)
     if del_method:
@@ -903,9 +903,12 @@ def generate_traverse_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -
 
 
 def generate_clear_for_class(cl: ClassIR, func_name: str, emitter: Emitter) -> None:
-    emitter.emit_line("static int")
-    emitter.emit_line(f"{func_name}({cl.struct_name(emitter.names)} *self)")
+    emitter.emit_line("static " + native_function_header(cl.clear, emitter))
     emitter.emit_line("{")
+    emitter.emit_line(
+        f"{cl.struct_name(emitter.names)} *self = "
+        f"({cl.struct_name(emitter.names)} *)cpy_r_self;"
+    )
     for base in reversed(cl.base_mro):
         for attr, rtype in base.attributes.items():
             emitter.emit_gc_clear(f"self->{emitter.attr(attr)}", rtype)
@@ -963,7 +966,7 @@ def generate_dealloc_for_class(
     if not cl.is_acyclic:
         emitter.emit_line("PyObject_GC_UnTrack(self);")
     if cl.builtin_base:
-        emitter.emit_line(f"{clear_func_name}(self);")
+        emitter.emit_line(f"{clear_func_name}((PyObject *)self);")
         # For native subclasses of builtins such as dict, the base deallocator
         # is responsible for tearing down base-owned storage and freeing memory.
         # Re-track self if base is GC-aware to match cpython's subtype_dealloc.
@@ -978,7 +981,7 @@ def generate_dealloc_for_class(
         emit_reuse_dealloc(cl, emitter)
     # The trashcan is needed to handle deep recursive deallocations
     emitter.emit_line(f"CPy_TRASHCAN_BEGIN(self, {dealloc_func_name})")
-    emitter.emit_line(f"{clear_func_name}(self);")
+    emitter.emit_line(f"{clear_func_name}((PyObject *)self);")
     emitter.emit_line("Py_TYPE(self)->tp_free((PyObject *)self);")
     emitter.emit_line("CPy_TRASHCAN_END(self)")
     emitter.emit_line("done: ;")

mypyc/ir/class_ir.py

@@ -7,7 +7,7 @@
 from mypyc.common import PROPSET_PREFIX, JsonDict
 from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
 from mypyc.ir.ops import DeserMaps, Value
-from mypyc.ir.rtypes import RInstance, RType, deserialize_type, object_rprimitive
+from mypyc.ir.rtypes import RInstance, RType, c_int_rprimitive, deserialize_type, object_rprimitive
 from mypyc.namegen import NameGenerator, exported_name
 
 # Some notes on the vtable layout: Each concrete class has a vtable
@@ -143,6 +143,13 @@ def __init__(
             module_name,
             FuncSignature([RuntimeArg("type", object_rprimitive)], RInstance(self)),
         )
+        self.clear = FuncDecl(
+            name + "_clear",
+            None,
+            module_name,
+            FuncSignature([RuntimeArg("self", RInstance(self))], c_int_rprimitive),
+            internal=True,
+        )
         # Attributes defined in the class (not inherited)
         self.attributes: dict[str, RType] = {}
         # Final attributes defined in the class (not inherited)
@@ -412,6 +419,7 @@ def serialize(self) -> JsonDict:
             "_serializable": self._serializable,
             "builtin_base": self.builtin_base,
             "ctor": self.ctor.serialize(),
+            "clear": self.clear.serialize(),
             # We serialize dicts as lists to ensure order is preserved
             "attributes": [(k, t.serialize()) for k, t in self.attributes.items()],
             "final_attributes": sorted(self.final_attributes),
@@ -474,6 +482,7 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> ClassIR:
         ir._serializable = data["_serializable"]
         ir.builtin_base = data["builtin_base"]
         ir.ctor = FuncDecl.deserialize(data["ctor"], ctx)
+        ir.clear = FuncDecl.deserialize(data["clear"], ctx)
         ir.attributes = {k: deserialize_type(t, ctx) for k, t in data["attributes"]}
         ir.final_attributes = set(data["final_attributes"])
         ir.method_decls = {

mypyc/irbuild/generator.py

@@ -23,6 +23,7 @@
     Call,
     Goto,
     Integer,
+    LoadErrorValue,
     MethodCall,
     RaiseStandardError,
     Register,
@@ -49,7 +50,7 @@
     load_outer_envs,
     setup_func_for_recursive_call,
 )
-from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl
+from mypyc.irbuild.nonlocalcontrol import ExceptNonlocalControl, gen_generator_func_cleanup
 from mypyc.irbuild.prepare import GENERATOR_HELPER_NAME
 from mypyc.primitives.exc_ops import (
     error_catch_op,
@@ -107,13 +108,20 @@ class that implements the function (each function gets a separate class).
         setup_func_for_recursive_call(
             builder, fitem, builder.fn_info.generator_class, prefix=GENERATOR_ATTRIBUTE_PREFIX
         )
+    cleanup_on_error = BasicBlock()
+    builder.builder.push_error_handler(cleanup_on_error)
     create_switch_for_generator_class(builder)
     add_raise_exception_blocks_to_generator_class(builder, fitem.line)
 
     add_vars_to_env(builder, prefix=GENERATOR_ATTRIBUTE_PREFIX)
 
     builder.accept(fitem.body)
     builder.maybe_add_implicit_return()
+    builder.builder.pop_error_handler()
+
+    builder.activate_block(cleanup_on_error)
+    gen_generator_func_cleanup(builder, fitem.line)
+    builder.add(Return(builder.add(LoadErrorValue(object_rprimitive))))
 
     populate_switch_for_generator_class(builder)
 

mypyc/irbuild/nonlocalcontrol.py

@@ -8,10 +8,13 @@
 from abc import abstractmethod
 from typing import TYPE_CHECKING
 
+from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.ops import (
+    ERR_NEVER,
     NO_TRACEBACK_LINE_NO,
     BasicBlock,
     Branch,
+    Call,
     Goto,
     Integer,
     Register,
@@ -90,10 +93,7 @@ class GeneratorNonlocalControl(BaseNonlocalControl):
     """Default nonlocal control in a generator function outside statements."""
 
     def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
-        # Assign an invalid next label number so that the next time
-        # __next__ is called, we jump to the case in which
-        # StopIteration is raised.
-        builder.assign(builder.fn_info.generator_class.next_label_target, Integer(-1), line)
+        gen_generator_func_cleanup(builder, line)
 
         # Raise a StopIteration containing a field for the value that
         # should be returned. Before doing so, create a new block
@@ -132,6 +132,25 @@ def gen_return(self, builder: IRBuilder, value: Value, line: int) -> None:
         builder.add(Return(Integer(0, object_rprimitive)))
 
 
+def gen_class_clear(builder: IRBuilder, obj: Value, cl: ClassIR, line: int) -> None:
+    call = Call(cl.clear, [obj], line)
+    call.error_kind = ERR_NEVER
+    builder.add(call)
+
+
+def gen_generator_func_cleanup(builder: IRBuilder, line: int) -> None:
+    """Clear references held by a completed generator or coroutine."""
+    cls = builder.fn_info.generator_class
+
+    # Assign an invalid next label number so that the next time __next__ is
+    # called, we jump to the case in which StopIteration is raised.
+    builder.assign(cls.next_label_target, Integer(-1), line)
+
+    gen_class_clear(builder, cls.curr_env_reg, builder.fn_info.env_class, line)
+    if builder.fn_info.env_class is not cls.ir:
+        gen_class_clear(builder, cls.self_reg, cls.ir, line)
+
+
 class CleanupNonlocalControl(NonlocalControl):
     """Abstract nonlocal control that runs some cleanup code."""
 

mypyc/test-data/run-async.test

@@ -407,6 +407,12 @@ import asyncio
 import gc
 import platform
 
+from testutil import is_gil_disabled
+
+def immediate_deallocation_checks_are_reliable() -> bool:
+    # Free-threaded builds can defer deallocation, so weakref liveness is not deterministic.
+    return not is_gil_disabled()
+
 async def sleep() -> None:
     # Non-zero value hangs on Windows.
     time = 0 if platform.system() == "Windows" else 0.0001
@@ -512,6 +518,106 @@ async def stolen(n: int) -> int:
 async def test_stolen() -> None:
     await assert_no_leaks(lambda: stolen(200), 16)
 
+async def retain_arg(c: set[int]) -> None:
+    async def nested() -> int:
+        await sleep()
+        return len(c)
+
+    await nested()
+
+async def test_completed_coroutine_releases_environment() -> None:
+    if not immediate_deallocation_checks_are_reliable():
+        return
+
+    import weakref
+
+    c = {1}
+    ref = weakref.ref(c)
+    coro = retain_arg(c)
+    c = {2}
+    await coro
+    assert ref() is None
+
+async def retain_arg_then_raise(c: set[int]) -> None:
+    async def nested() -> int:
+        await sleep()
+        return len(c)
+
+    await nested()
+    raise ValueError
+
+async def test_failed_coroutine_releases_environment() -> None:
+    if not immediate_deallocation_checks_are_reliable():
+        return
+
+    import weakref
+
+    c = {1}
+    ref = weakref.ref(c)
+    coro = retain_arg_then_raise(c)
+    c = {2}
+    try:
+        await coro
+    except ValueError:
+        pass
+    assert ref() is None
+
+async def close_retain_arg(c: set[int]) -> None:
+    await sleep()
+    len(c)
+
+async def test_closed_coroutine_releases_environment() -> None:
+    if not immediate_deallocation_checks_are_reliable():
+        return
+
+    import weakref
+    from typing import Any, cast
+
+    c = {1}
+    ref = weakref.ref(c)
+    coro = cast(Any, close_retain_arg(c))
+    coro.send(None)
+    c = {2}
+    coro.close()
+    assert ref() is None
+
+async def interpreted_leaf(c: set[int]) -> int:
+    await sleep()
+    return len(c)
+
+async def compiled_leaf(c: set[int]) -> int:
+    await sleep()
+    return len(c)
+
+async def compiled_awaits_interpreted(c: set[int]) -> int:
+    import interpreted
+
+    return await interpreted.interpreted_leaf(c)
+
+async def interpreted_awaits_compiled(fn, c: set[int]) -> int:
+    return await fn(c)
+
+async def test_completed_mixed_coroutine_releases_environment() -> None:
+    if not immediate_deallocation_checks_are_reliable():
+        return
+
+    import interpreted
+    import weakref
+
+    c = {1}
+    ref = weakref.ref(c)
+    coro = compiled_awaits_interpreted(c)
+    c = {2}
+    assert await coro == 1
+    assert ref() is None
+
+    c = {3}
+    ref = weakref.ref(c)
+    coro = interpreted.interpreted_awaits_compiled(compiled_leaf, c)
+    c = {4}
+    assert await coro == 1
+    assert ref() is None
+
 [file asyncio/__init__.pyi]
 async def sleep(t: float) -> None: ...