gh-131527: Stackref debug borrow checker (#140599)

Add borrow checking to the stackref debug mode

---------

Co-authored-by: mpage <mpage@meta.com>

Include/internal/pycore_stackref.h

@@ -63,6 +63,8 @@ PyAPI_FUNC(PyObject *) _Py_stackref_get_object(_PyStackRef ref);
 PyAPI_FUNC(PyObject *) _Py_stackref_close(_PyStackRef ref, const char *filename, int linenumber);
 PyAPI_FUNC(_PyStackRef) _Py_stackref_create(PyObject *obj, uint16_t flags, const char *filename, int linenumber);
 PyAPI_FUNC(void) _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber);
+PyAPI_FUNC(_PyStackRef) _Py_stackref_get_borrowed_from(_PyStackRef ref, const char *filename, int linenumber);
+PyAPI_FUNC(void) _Py_stackref_set_borrowed_from(_PyStackRef ref, _PyStackRef borrowed_from, const char *filename, int linenumber);
 extern void _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref);
 
 static const _PyStackRef PyStackRef_NULL = { .index = 0 };
@@ -248,7 +250,12 @@ _PyStackRef_DUP(_PyStackRef ref, const char *filename, int linenumber)
     } else {
         flags = Py_TAG_REFCNT;
     }
-    return _Py_stackref_create(obj, flags, filename, linenumber);
+    _PyStackRef new_ref = _Py_stackref_create(obj, flags, filename, linenumber);
+    if (flags == Py_TAG_REFCNT && !_Py_IsImmortal(obj)) {
+        _PyStackRef borrowed_from = _Py_stackref_get_borrowed_from(ref, filename, linenumber);
+        _Py_stackref_set_borrowed_from(new_ref, borrowed_from, filename, linenumber);
+    }
+    return new_ref;
 }
 #define PyStackRef_DUP(REF) _PyStackRef_DUP(REF, __FILE__, __LINE__)
 
@@ -259,6 +266,7 @@ _PyStackRef_CLOSE_SPECIALIZED(_PyStackRef ref, destructor destruct, const char *
     assert(!PyStackRef_IsNull(ref));
     assert(!PyStackRef_IsTaggedInt(ref));
     PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
+    assert(Py_REFCNT(obj) > 0);
     if (PyStackRef_RefcountOnObject(ref)) {
         _Py_DECREF_SPECIALIZED(obj, destruct);
     }
@@ -274,7 +282,11 @@ _PyStackRef_Borrow(_PyStackRef ref, const char *filename, int linenumber)
         return ref;
     }
     PyObject *obj = _Py_stackref_get_object(ref);
-    return _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
+    _PyStackRef new_ref = _Py_stackref_create(obj, Py_TAG_REFCNT, filename, linenumber);
+    if (!_Py_IsImmortal(obj)) {
+        _Py_stackref_set_borrowed_from(new_ref, ref, filename, linenumber);
+    }
+    return new_ref;
 }
 #define PyStackRef_Borrow(REF) _PyStackRef_Borrow((REF), __FILE__, __LINE__)
 
@@ -310,13 +322,22 @@ PyStackRef_IsHeapSafe(_PyStackRef ref)
 static inline _PyStackRef
 _PyStackRef_MakeHeapSafe(_PyStackRef ref, const char *filename, int linenumber)
 {
-    if (PyStackRef_IsHeapSafe(ref)) {
+    // Special references that can't be closed.
+    if (ref.index < INITIAL_STACKREF_INDEX) {
         return ref;
     }
 
+    bool heap_safe = PyStackRef_IsHeapSafe(ref);
     PyObject *obj = _Py_stackref_close(ref, filename, linenumber);
-    Py_INCREF(obj);
+    uint16_t flags = 0;
-    return _Py_stackref_create(obj, 0, filename, linenumber);
+    if (heap_safe) {
+        // Close old ref and create a new one with the same flags.
+        // This is necessary for correct borrow checking.
+        flags = ref.index & Py_TAG_BITS;
+    } else {
+        Py_INCREF(obj);
+    }
+    return _Py_stackref_create(obj, flags, filename, linenumber);
 }
 #define PyStackRef_MakeHeapSafe(REF) _PyStackRef_MakeHeapSafe(REF, __FILE__, __LINE__)
 

Misc/NEWS.d/next/Core_and_Builtins/2025-10-25-21-31-43.gh-issue-131527.V-JVNP.rst

@@ -0,0 +1,2 @@
+Dynamic borrow checking for stackrefs is added to ``Py_STACKREF_DEBUG``
+mode. Patch by Mikhail Efimov.

Python/stackrefs.c

@@ -19,6 +19,8 @@ typedef struct _table_entry {
     int linenumber;
     const char *filename_borrow;
     int linenumber_borrow;
+    int borrows;
+    _PyStackRef borrowed_from;
 } TableEntry;
 
 TableEntry *
@@ -34,6 +36,8 @@ make_table_entry(PyObject *obj, const char *filename, int linenumber)
     result->linenumber = linenumber;
     result->filename_borrow = NULL;
     result->linenumber_borrow = 0;
+    result->borrows = 0;
+    result->borrowed_from = PyStackRef_NULL;
     return result;
 }
 
@@ -47,11 +51,13 @@ _Py_stackref_get_object(_PyStackRef ref)
     PyInterpreterState *interp = PyInterpreterState_Get();
     assert(interp != NULL);
     if (ref.index >= interp->next_stackref) {
-        _Py_FatalErrorFormat(__func__, "Garbled stack ref with ID %" PRIu64 "\n", ref.index);
+        _Py_FatalErrorFormat(__func__,
+            "Garbled stack ref with ID %" PRIu64 "\n", ref.index);
     }
     TableEntry *entry = _Py_hashtable_get(interp->open_stackrefs_table, (void *)ref.index);
     if (entry == NULL) {
-        _Py_FatalErrorFormat(__func__, "Accessing closed stack ref with ID %" PRIu64 "\n", ref.index);
+        _Py_FatalErrorFormat(__func__,
+            "Accessing closed stack ref with ID %" PRIu64 "\n", ref.index);
     }
     return entry->obj;
 }
@@ -68,13 +74,16 @@ _Py_stackref_close(_PyStackRef ref, const char *filename, int linenumber)
     assert(!PyStackRef_IsError(ref));
     PyInterpreterState *interp = PyInterpreterState_Get();
     if (ref.index >= interp->next_stackref) {
-        _Py_FatalErrorFormat(__func__, "Invalid StackRef with ID %" PRIu64 " at %s:%d\n", (void *)ref.index, filename, linenumber);
+        _Py_FatalErrorFormat(__func__,
-
+            "Invalid StackRef with ID %" PRIu64 " at %s:%d\n",
+            ref.index, filename, linenumber);
     }
     PyObject *obj;
     if (ref.index < INITIAL_STACKREF_INDEX) {
         if (ref.index == 0) {
-            _Py_FatalErrorFormat(__func__, "Passing NULL to PyStackRef_CLOSE at %s:%d\n", filename, linenumber);
+            _Py_FatalErrorFormat(__func__,
+                "Passing NULL to _Py_stackref_close at %s:%d\n",
+                filename, linenumber);
         }
         // Pre-allocated reference to None, False or True -- Do not clear
         TableEntry *entry = _Py_hashtable_get(interp->open_stackrefs_table, (void *)ref.index);
@@ -88,10 +97,27 @@ _Py_stackref_close(_PyStackRef ref, const char *filename, int linenumber)
             if (entry != NULL) {
                 _Py_FatalErrorFormat(__func__,
                     "Double close of ref ID %" PRIu64 " at %s:%d. Referred to instance of %s at %p. Closed at %s:%d\n",
-                    (void *)ref.index, filename, linenumber, entry->classname, entry->obj, entry->filename, entry->linenumber);
+                    ref.index, filename, linenumber, entry->classname, entry->obj, entry->filename, entry->linenumber);
             }
 #endif
-            _Py_FatalErrorFormat(__func__, "Invalid StackRef with ID %" PRIu64 "\n", (void *)ref.index);
+            _Py_FatalErrorFormat(__func__,
+                "Invalid StackRef with ID %" PRIu64 " at %s:%d\n",
+                ref.index, filename, linenumber);
+        }
+        if (!PyStackRef_IsNull(entry->borrowed_from)) {
+            _PyStackRef borrowed_from = entry->borrowed_from;
+            TableEntry *entry_borrowed = _Py_hashtable_get(interp->open_stackrefs_table, (void *)borrowed_from.index);
+            if (entry_borrowed == NULL) {
+                _Py_FatalErrorFormat(__func__,
+                    "Invalid borrowed StackRef with ID %" PRIu64 " at %s:%d\n",
+                    borrowed_from.index, filename, linenumber);
+            }
+            entry_borrowed->borrows--;
+        }
+        if (entry->borrows > 0) {
+            _Py_FatalErrorFormat(__func__,
+                "StackRef with ID %" PRIu64 " closed with %d borrowed refs at %s:%d. Opened at %s:%d\n",
+                ref.index, entry->borrows, filename, linenumber, entry->filename, entry->linenumber);
         }
         obj = entry->obj;
         free(entry);
@@ -143,15 +169,62 @@ _Py_stackref_record_borrow(_PyStackRef ref, const char *filename, int linenumber
         if (entry != NULL) {
             _Py_FatalErrorFormat(__func__,
                 "Borrow of closed ref ID %" PRIu64 " at %s:%d. Referred to instance of %s at %p. Closed at %s:%d\n",
-                (void *)ref.index, filename, linenumber, entry->classname, entry->obj, entry->filename, entry->linenumber);
+                ref.index, filename, linenumber, entry->classname, entry->obj, entry->filename, entry->linenumber);
         }
 #endif
-        _Py_FatalErrorFormat(__func__, "Invalid StackRef with ID %" PRIu64 " at %s:%d\n", (void *)ref.index, filename, linenumber);
+        _Py_FatalErrorFormat(__func__,
+            "Invalid StackRef with ID %" PRIu64 " at %s:%d\n",
+            ref.index, filename, linenumber);
     }
     entry->filename_borrow = filename;
     entry->linenumber_borrow = linenumber;
 }
 
+_PyStackRef
+_Py_stackref_get_borrowed_from(_PyStackRef ref, const char *filename, int linenumber)
+{
+    assert(!PyStackRef_IsError(ref));
+    PyInterpreterState *interp = PyInterpreterState_Get();
+
+    TableEntry *entry = _Py_hashtable_get(interp->open_stackrefs_table, (void *)ref.index);
+    if (entry == NULL) {
+        _Py_FatalErrorFormat(__func__,
+            "Invalid StackRef with ID %" PRIu64 " at %s:%d\n",
+            ref.index, filename, linenumber);
+    }
+
+    return entry->borrowed_from;
+}
+
+// This function should be used no more than once per ref.
+void
+_Py_stackref_set_borrowed_from(_PyStackRef ref, _PyStackRef borrowed_from, const char *filename, int linenumber)
+{
+    assert(!PyStackRef_IsError(ref));
+    PyInterpreterState *interp = PyInterpreterState_Get();
+
+    TableEntry *entry = _Py_hashtable_get(interp->open_stackrefs_table, (void *)ref.index);
+    if (entry == NULL) {
+        _Py_FatalErrorFormat(__func__,
+            "Invalid StackRef (ref) with ID %" PRIu64 " at %s:%d\n",
+            ref.index, filename, linenumber);
+    }
+
+    assert(PyStackRef_IsNull(entry->borrowed_from));
+    if (PyStackRef_IsNull(borrowed_from)) {
+        return;
+    }
+
+    TableEntry *entry_borrowed = _Py_hashtable_get(interp->open_stackrefs_table, (void *)borrowed_from.index);
+    if (entry_borrowed == NULL) {
+        _Py_FatalErrorFormat(__func__,
+            "Invalid StackRef (borrowed_from) with ID %" PRIu64 " at %s:%d\n",
+            borrowed_from.index, filename, linenumber);
+    }
+
+    entry->borrowed_from = borrowed_from;
+    entry_borrowed->borrows++;
+}
 
 void
 _Py_stackref_associate(PyInterpreterState *interp, PyObject *obj, _PyStackRef ref)