[mypyc] Fix function wrapper memory leak (#21654)

The `name` and `code` variables passed to the function wrapper init
function are increfed inside it before storing them in the function
wrapper object but they are not decrefed after the init function returns
which creates a leak.

Change the init function to instead steal `name` and `code` and document
it.

Author: p-sawicki

mypyc/lib-rt/function_wrapper.c

@@ -196,6 +196,7 @@ static PyObject* CPyFunction_Vectorcall(PyObject *func, PyObject *const *args, s
 }
 
 
+// Steals ml, name, and code. Borrows module.
 static CPyFunction* CPyFunction_Init(CPyFunction *op, PyMethodDef *ml, PyObject* name,
                                      PyObject *module, PyObject* code, bool set_self) {
     PyCFunctionObject *cf = (PyCFunctionObject *)op;
@@ -206,12 +207,10 @@ static CPyFunction* CPyFunction_Init(CPyFunction *op, PyMethodDef *ml, PyObject*
     Py_XINCREF(module);
     cf->m_module = module;
 
-    Py_INCREF(name);
     op->func_name = name;
 
     ((PyCMethodObject *)op)->mm_class = NULL;
 
-    Py_XINCREF(code);
     op->func_code = code;
 
     CPyFunction_func_vectorcall(op) = CPyFunction_Vectorcall;
@@ -243,7 +242,7 @@ PyObject* CPyFunction_New(PyObject *module, const char *filename, const char *fu
                           PyCFunction func, int func_flags, const char *func_doc,
                           int first_line, int code_flags, bool has_self_arg) {
     PyMethodDef *method = NULL;
-    PyObject *code = NULL, *op = NULL;
+    PyObject *code = NULL, *name = NULL, *op = NULL;
     bool set_self = false;
 
 #ifdef Py_GIL_DISABLED
@@ -280,18 +279,24 @@ PyObject* CPyFunction_New(PyObject *module, const char *filename, const char *fu
     if (unlikely(!code)) {
         goto err;
     }
+    name = PyUnicode_FromString(funcname);
+    if (unlikely(!name)) {
+        goto err;
+    }
 
     // Set m_self inside the function wrapper only if the wrapped function has no self arg
     // to pass m_self as the self arg when the function is called.
     // When the function has a self arg, it will come in the args vector passed to the
     // vectorcall handler.
     set_self = !has_self_arg;
-    op = (PyObject *)CPyFunction_Init(PyObject_GC_New(CPyFunction, CPyFunctionType),
-                                      method, PyUnicode_FromString(funcname), module,
-                                      code, set_self);
-    if (unlikely(!op)) {
+    CPyFunction *raw = PyObject_GC_New(CPyFunction, CPyFunctionType);
+    if (unlikely(!raw)) {
         goto err;
     }
+    op = (PyObject *)CPyFunction_Init(raw, method, name, module, code, set_self);
+    method = NULL;
+    name = NULL;
+    code = NULL;
     PyObject_GC_Track(op);
     return op;
 
@@ -300,5 +305,7 @@ PyObject* CPyFunction_New(PyObject *module, const char *filename, const char *fu
     if (method) {
         PyMem_Free(method);
     }
+    Py_XDECREF(name);
+    Py_XDECREF(code);
     return NULL;
 }

mypyc/test-data/run-async.test

@@ -1416,13 +1416,16 @@ def run(x: object) -> object: ...
 
 [case testAsyncIntrospection]
 import asyncio
+import gc
 import inspect
 import sys
 import weakref
 
 from functools import wraps
 from typing import Any, Callable, TypeVar, cast
 
+from testutil import is_gil_disabled
+
 def identity(val: int) -> int:
     return val
 
@@ -1589,6 +1592,37 @@ def test_nested() -> None:
     assert is_coroutine(nested_wrapped_async)
     assert asyncio.run(nested_wrapped_async()) == 4
 
+def test_async_function_wrapper_code_refcount() -> None:
+    if is_gil_disabled():
+        # On free-threaded builds the code object might be immortal, so the ref count test doesn't work.
+        return
+    code = getattr(identity_async, "__code__")
+    getrefcount = getattr(sys, "getrefcount")
+    # getrefcount sees the local code variable plus the wrapper-owned reference.
+    assert getrefcount(code) == 2, getrefcount(code)
+
+def test_nested_async_function_wrapper_code_refcount() -> None:
+    if is_gil_disabled():
+        # On free-threaded builds the code object might be immortal, so the ref count test doesn't work.
+        return
+    def make_nested() -> Any:
+        async def nested_refcounted() -> int:
+            return 1
+
+        return nested_refcounted
+
+    getrefcount = getattr(sys, "getrefcount")
+    fn = make_nested()
+    code = getattr(fn, "__code__")
+    before = getrefcount(code)
+    assert asyncio.run(fn()) == 1
+
+    del fn
+    gc.collect()
+    after = getrefcount(code)
+    assert before == after + 1, (before, after)
+    assert after == 1, after
+
 [file asyncio/__init__.pyi]
 def run(x: object) -> object: ...