[mypyc] Create function wrappers of nested coroutines lazily (#22048)

Each nested async def instance used to get it's own freshly allocated
function wrapper object, which is only needed relatively rarely for
introspection (e.g. `__name__`). This was slow to construct. Now we
construct this lazily. On free-threaded builds we need to deal with
other threads racing as well.

Top-level coroutines and methods are unaffected. Their wrappers were
already created once at import time.

Microbenchmark results (Python 3.14, 1M iterations):
* create a nested coroutine closure: 197 ns -> 16 ns
* create it and run the coroutine: 713 ns -> 526 ns

Author: JukkaL

mypyc/codegen/emitclass.py

@@ -1498,7 +1498,20 @@ def success() -> None:
         wrapper_name = emit_instance(cl.methods["__call__"], cl.coroutine_name)
         struct_name = cl.struct_name(emitter.names)
         attr = emitter.attr(CPYFUNCTION_NAME)
-        emitter.emit_line(f"(({struct_name} *)type)->{attr} = {wrapper_name};")
+        if IS_FREE_THREADED:
+            emitter.emit_line(
+                f"CPy_InitAttrRefIfNull(type, (PyObject **)&(({struct_name} *)type)->{attr}, "
+                f"{wrapper_name});"
+            )
+        else:
+            # Creating the wrapper may run arbitrary code (e.g. via GC), which could
+            # have initialized the attribute already.
+            field = f"(({struct_name} *)type)->{attr}"
+            emitter.emit_line(f"if ({field} == NULL) {{")
+            emitter.emit_line(f"{field} = {wrapper_name};")
+            emitter.emit_line("} else {")
+            emitter.emit_line(f"Py_DECREF({wrapper_name});")
+            emitter.emit_line("}")
         return success()
 
     if not any(fn.decl.is_coroutine for fn in cl.methods.values()):

mypyc/irbuild/callable_class.py

@@ -6,10 +6,26 @@
 
 from __future__ import annotations
 
-from mypyc.common import CPYFUNCTION_NAME, ENV_ATTR_NAME, PROPSET_PREFIX, SELF_NAME
+from mypyc.common import (
+    CPYFUNCTION_NAME,
+    ENV_ATTR_NAME,
+    IS_FREE_THREADED,
+    PROPSET_PREFIX,
+    SELF_NAME,
+)
 from mypyc.ir.class_ir import ClassIR
 from mypyc.ir.func_ir import FuncDecl, FuncIR, FuncSignature, RuntimeArg
-from mypyc.ir.ops import BasicBlock, Call, GetAttr, Integer, Register, Return, SetAttr, Value
+from mypyc.ir.ops import (
+    BasicBlock,
+    Branch,
+    Call,
+    GetAttr,
+    Integer,
+    Register,
+    Return,
+    SetAttr,
+    Value,
+)
 from mypyc.ir.rtypes import RInstance, c_pointer_rprimitive, int_rprimitive, object_rprimitive
 from mypyc.irbuild.builder import IRBuilder
 from mypyc.irbuild.context import FuncInfo, ImplicitClass
@@ -119,7 +135,24 @@ def add_coroutine_properties(
     line = builder.fn_info.fitem.line
 
     def get_func_wrapper() -> Value:
-        return builder.add(GetAttr(builder.self(), CPYFUNCTION_NAME, line))
+        # The wrapper is created lazily on first introspection, since creating it
+        # (including a code object) for every instance is expensive.
+        self_reg = builder.self()
+        init, done = BasicBlock(), BasicBlock()
+        cur = builder.add(
+            GetAttr(
+                self_reg,
+                CPYFUNCTION_NAME,
+                line,
+                borrow=not IS_FREE_THREADED,
+                allow_error_value=True,
+            )
+        )
+        builder.add(Branch(cur, init, done, Branch.IS_ERROR))
+        builder.activate_block(init)
+        builder.add_coroutine_setup_call(callable_class_ir.name, self_reg)
+        builder.goto_and_activate(done)
+        return builder.add(GetAttr(self_reg, CPYFUNCTION_NAME, line))
 
     for name, primitive in properties.items():
         with builder.enter_method(callable_class_ir, name, object_rprimitive, internal=True):
@@ -243,9 +276,6 @@ def instantiate_callable_class(builder: IRBuilder, fn_info: FuncInfo) -> Value:
         # happens before the callable can escape.
         set_env.mark_as_initializer()
         builder.add(set_env)
-    # Initialize function wrapper for callable classes. As opposed to regular functions,
-    # each instance of a callable class needs its own wrapper because they might be instantiated
-    # inside other functions.
-    if not fn_info.in_non_ext and fn_info.is_coroutine:
-        builder.add_coroutine_setup_call(fn_info.callable_class.ir.name, func_reg)
+    # Callable classes of coroutines need their own function wrapper per instance
+    # (for introspection), but it's created lazily by the property getters/setters.
     return func_reg

mypyc/lib-rt/pythonsupport.h

@@ -129,6 +129,19 @@ static inline void CPy_SetAttrRef(PyObject *owner, PyObject **field, PyObject *v
     Py_XDECREF(old);
 }
 
+// Publish a lazily created attribute once. Concurrent initializers may create
+// separate values; only the first is installed, and the others are discarded.
+// Decref outside the owner's critical section in case it runs Python code.
+static inline void CPy_InitAttrRefIfNull(PyObject *owner, PyObject **field, PyObject *value) {
+    Py_BEGIN_CRITICAL_SECTION(owner);
+    if (_Py_atomic_load_ptr_relaxed(field) == NULL) {
+        _Py_atomic_store_ptr_release(field, value);
+        value = NULL;
+    }
+    Py_END_CRITICAL_SECTION();
+    Py_XDECREF(value);
+}
+
 // Initialize a native attribute that is known to be previously undefined (NULL),
 // stealing the reference to 'value'.
 //

mypyc/test-data/run-async.test

@@ -2162,6 +2162,98 @@ from typing import Any, Generator
 
 def run(x: object) -> object: ...
 
+[case testNestedCoroutineLazyFunctionWrapper]
+import asyncio
+import inspect
+from typing import Any
+
+def make(n: int) -> Any:
+    async def inner(x: int = 5) -> int:
+        return x + n
+
+    return inner
+
+def test_set_name_first() -> None:
+    f = make(1)
+    # The first introspection access is a setter.
+    f.__name__ = "renamed"
+    assert f.__name__ == "renamed"
+    assert inspect.iscoroutinefunction(f)
+    assert asyncio.run(f()) == 6
+
+def test_instances_are_independent() -> None:
+    a = make(1)
+    b = make(2)
+    assert asyncio.run(a(1)) == 2
+    a.__name__ = "a"
+    assert b.__name__ == "inner"
+    assert a.__name__ == "a"
+    assert asyncio.run(b(1)) == 3
+
+def test_freed_instance_reuse() -> None:
+    for i in range(5):
+        f = make(i)
+        assert f.__name__ == "inner"
+        f.__name__ = f"name{i}"
+        assert f.__code__ is not None
+        del f
+        # Never introspected, so the wrapper is never created.
+        g = make(i)
+        assert asyncio.run(g(0)) == i
+        del g
+
+[file asyncio/__init__.pyi]
+def run(x: object) -> object: ...
+
+[case testNestedCoroutineConcurrentFirstIntrospection]
+from typing import Any
+
+def make() -> Any:
+    async def inner() -> None:
+        pass
+
+    return inner
+
+[file driver.py]
+import gc
+import threading
+
+from native import make
+
+workers = 16
+rounds = 300
+start = threading.Barrier(workers + 1)
+done = threading.Barrier(workers + 1)
+current = [None]
+
+def worker():
+    for _ in range(rounds):
+        start.wait()
+        f = current[0]
+        assert f.__name__ == "inner"
+        done.wait()
+
+threads = [threading.Thread(target=worker) for _ in range(workers)]
+for thread in threads:
+    thread.start()
+
+f = make()
+assert f.__name__ == "inner"
+wrapper_type = type(f.__cpyfunction__)
+del f
+gc.collect()
+before = sum(type(obj) is wrapper_type for obj in gc.get_objects())
+for _ in range(rounds):
+    current[0] = make()
+    start.wait()
+    done.wait()
+    current[0] = None
+for thread in threads:
+    thread.join()
+gc.collect()
+after = sum(type(obj) is wrapper_type for obj in gc.get_objects())
+assert after == before, (before, after)
+
 [case testCPyFunctionVariableNameConflict]
 class C:
     async def operation(self) -> int: