Fix stackref leaks

Author: markshannon

Include/internal/pycore_stackref.h

@@ -66,6 +66,7 @@ typedef union _PyStackRef {
 PyAPI_FUNC(PyObject *) _Py_stackref_get_object(_PyStackRef ref);
 PyAPI_FUNC(PyObject *) _Py_stackref_close(_PyStackRef ref);
 PyAPI_FUNC(_PyStackRef) _Py_stackref_create(PyObject *obj, const char *filename, int linenumber);
+PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber);
 extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);
 
 static const _PyStackRef PyStackRef_NULL = { .index = 0 };
@@ -99,11 +100,14 @@ PyStackRef_IsNone(_PyStackRef ref)
 }
 
 static inline PyObject *
-PyStackRef_AsPyObjectBorrow(_PyStackRef ref)
+_PyStackRef_AsPyObjectBorrow(_PyStackRef ref, const char *filename, int linenumber)
 {
+    _Py_stackref_record_borrow(ref, filename, linenumber);
     return _Py_stackref_get_object(ref);
 }
 
+#define PyStackRef_AsPyObjectBorrow(REF) _PyStackRef_AsPyObjectBorrow((REF), __FILE__, __LINE__)
+
 static inline PyObject *
 PyStackRef_AsPyObjectSteal(_PyStackRef ref)
 {

Python/bytecodes.c

@@ -679,7 +679,7 @@ dummy_func(
             assert(Py_REFCNT(left_o) >= 2);
             PyStackRef_CLOSE(left);
             DEAD(left);
-            PyObject *temp = PyStackRef_AsPyObjectBorrow(*target_local);
+            PyObject *temp = PyStackRef_AsPyObjectSteal(*target_local);
             PyUnicode_Append(&temp, right_o);
             *target_local = PyStackRef_FromPyObjectSteal(temp);
             PyStackRef_CLOSE_SPECIALIZED(right, _PyUnicode_ExactDealloc);
@@ -4476,17 +4476,17 @@ dummy_func(
 
         op(_DO_CALL_FUNCTION_EX, (func_st, unused, callargs_st, kwargs_st if (oparg & 1) -- result)) {
             PyObject *func = PyStackRef_AsPyObjectBorrow(func_st);
-            PyObject *callargs = PyStackRef_AsPyObjectBorrow(callargs_st);
-            PyObject *kwargs = PyStackRef_AsPyObjectBorrow(kwargs_st);
 
             // DICT_MERGE is called before this opcode if there are kwargs.
             // It converts all dict subtypes in kwargs into regular dicts.
-            assert(kwargs == NULL || PyDict_CheckExact(kwargs));
-            assert(PyTuple_CheckExact(callargs));
             EVAL_CALL_STAT_INC_IF_FUNCTION(EVAL_CALL_FUNCTION_EX, func);
             PyObject *result_o;
             assert(!_PyErr_Occurred(tstate));
             if (opcode == INSTRUMENTED_CALL_FUNCTION_EX) {
+                PyObject *callargs = PyStackRef_AsPyObjectBorrow(callargs_st);
+                PyObject *kwargs = PyStackRef_AsPyObjectBorrow(kwargs_st);
+                assert(kwargs == NULL || PyDict_CheckExact(kwargs));
+                assert(PyTuple_CheckExact(callargs));
                 PyObject *arg = PyTuple_GET_SIZE(callargs) > 0 ?
                     PyTuple_GET_ITEM(callargs, 0) : &_PyInstrumentation_MISSING;
                 int err = _Py_call_instrumentation_2args(
@@ -4517,7 +4517,10 @@ dummy_func(
                 if (Py_TYPE(func) == &PyFunction_Type &&
                     tstate->interp->eval_frame == NULL &&
                     ((PyFunctionObject *)func)->vectorcall == _PyFunction_Vectorcall) {
+                    PyObject *callargs = PyStackRef_AsPyObjectSteal(callargs_st);
                     assert(PyTuple_CheckExact(callargs));
+                    PyObject *kwargs = PyStackRef_IsNull(kwargs_st) ? NULL : PyStackRef_AsPyObjectSteal(kwargs_st);
+                    assert(kwargs == NULL || PyDict_CheckExact(kwargs));
                     Py_ssize_t nargs = PyTuple_GET_SIZE(callargs);
                     int code_flags = ((PyCodeObject *)PyFunction_GET_CODE(func))->co_flags;
                     PyObject *locals = code_flags & CO_OPTIMIZED ? NULL : Py_NewRef(PyFunction_GET_GLOBALS(func));
@@ -4535,6 +4538,10 @@ dummy_func(
                     frame->return_offset = 1;
                     DISPATCH_INLINED(new_frame);
                 }
+                PyObject *callargs = PyStackRef_AsPyObjectBorrow(callargs_st);
+                assert(PyTuple_CheckExact(callargs));
+                PyObject *kwargs = PyStackRef_AsPyObjectBorrow(kwargs_st);
+                assert(kwargs == NULL || PyDict_CheckExact(kwargs));
                 result_o = PyObject_Call(func, callargs, kwargs);
             }
             PyStackRef_XCLOSE(kwargs_st);

Python/executor_cases.c.h

@@ -17,6 +17,8 @@ typedef struct _table_entry {
     const char *classname;
     const char *filename;
     int linenumber;
+    const char *filename_borrow;
+    int linenumber_borrow;
 } TableEntry;
 
 TableEntry *
@@ -30,6 +32,7 @@ make_table_entry(PyObject *obj, const char *filename, int linenumber)
     result->classname = Py_TYPE(obj)->tp_name;
     result->filename = filename;
     result->linenumber = linenumber;
+    result->filename_borrow = NULL;
     return result;
 }
 
@@ -92,6 +95,22 @@ _Py_stackref_create(PyObject *obj, const char *filename, int linenumber)
     return (_PyStackRef){ .index = new_id };
 }
 
+void
+_Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber)
+{
+    if (ref.index <= 3) {
+        return;
+    }
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    TableEntry *entry = _Py_hashtable_get(interp->stackref_debug_table, (void *)ref.index);
+    if (entry == NULL) {
+        _Py_FatalErrorFormat(__func__, "Invalid StackRef with ID %" PRIu64 "\n", (void *)ref.index);
+    }
+    entry->filename_borrow = filename;
+    entry->linenumber_borrow = linenumber;
+}
+
+
 void
 _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref)
 {
@@ -108,20 +127,29 @@ _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef re
 
 
 static int
-report_leak(_Py_hashtable_t *ht, const void *key, const void *value, void *user_data)
+report_leak(_Py_hashtable_t *ht, const void *key, const void *value, void *leak)
 {
     TableEntry *entry = (TableEntry *)value;
     if (!_Py_IsStaticImmortal(entry->obj)) {
-        printf("Stackref leak. Refers to instance of %s at %p. Created at %s:%d\n",
+        *(int *)leak = 1;
+        printf("Stackref leak. Refers to instance of %s at %p. Created at %s:%d",
                entry->classname, entry->obj, entry->filename, entry->linenumber);
+        if (entry->filename_borrow != NULL) {
+            printf(". Last borrow at %s:%d",entry->filename_borrow, entry->linenumber_borrow);
+        }
+        printf("\n");
     }
     return 0;
 }
 
 void
 _Py_stackref_report_leaks(PyInterpreterState *interp)
 {
-    _Py_hashtable_foreach(interp->stackref_debug_table, report_leak, NULL);
+    int leak = 0;
+    _Py_hashtable_foreach(interp->stackref_debug_table, report_leak, &leak);
+    if (leak) {
+        Py_FatalError("Stackrefs leaked.");
+    }
 }
 
 #endif