gh-131527: Stackref debug borrow checker

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);
- return _Py_stackref_create(obj, 0, filename, linenumber);
+ uint16_t flags = 0;
+ 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)