gh-124218: Use per-thread refcounts for code objects (#125216)

Use per-thread refcounting for the reference from function objects to
their corresponding code object. This can be a source of contention when
frequently creating nested functions. Deferred refcounting alone isn't a
great fit here because these references are on the heap and may be
modified by other libraries.

2 files changed, 56 insertions, 34 deletions

diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
index 8832692..ad1a7d7 100644
--- a/Include/internal/pycore_object.h
+++ b/Include/internal/pycore_object.h
@@ -14,7 +14,7 @@ extern "C" {
 #include "pycore_interp.h"        // PyInterpreterState.gc
 #include "pycore_pyatomic_ft_wrappers.h"  // FT_ATOMIC_STORE_PTR_RELAXED
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
-#include "pycore_uniqueid.h"      // _PyType_IncrefSlow
+#include "pycore_uniqueid.h"      // _PyObject_ThreadIncrefSlow()
 
 // This value is added to `ob_ref_shared` for objects that use deferred
 // reference counting so that they are not immediately deallocated when the
@@ -291,8 +291,32 @@ extern bool _PyRefchain_IsTraced(PyInterpreterState *interp, PyObject *obj);
 #ifndef Py_GIL_DISABLED
 #  define _Py_INCREF_TYPE Py_INCREF
 #  define _Py_DECREF_TYPE Py_DECREF
+#  define _Py_INCREF_CODE Py_INCREF
+#  define _Py_DECREF_CODE Py_DECREF
 #else
 static inline void
+_Py_THREAD_INCREF_OBJECT(PyObject *obj, Py_ssize_t unique_id)
+{
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
+
+    // Unsigned comparison so that `unique_id=-1`, which indicates that
+    // per-thread refcounting has been disabled on this object, is handled by
+    // the "else".
+    if ((size_t)unique_id < (size_t)tstate->refcounts.size) {
+#  ifdef Py_REF_DEBUG
+        _Py_INCREF_IncRefTotal();
+#  endif
+        _Py_INCREF_STAT_INC();
+        tstate->refcounts.values[unique_id]++;
+    }
+    else {
+        // The slow path resizes the per-thread refcount array if necessary.
+        // It handles the unique_id=-1 case to keep the inlinable function smaller.
+        _PyObject_ThreadIncrefSlow(obj, unique_id);
+    }
+}
+
+static inline void
 _Py_INCREF_TYPE(PyTypeObject *type)
 {
     if (!_PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
@@ -308,29 +332,38 @@ _Py_INCREF_TYPE(PyTypeObject *type)
 #  pragma GCC diagnostic push
 #  pragma GCC diagnostic ignored "-Warray-bounds"
 #endif
+    _Py_THREAD_INCREF_OBJECT((PyObject *)type, ((PyHeapTypeObject *)type)->unique_id);
+#if defined(__GNUC__) && __GNUC__ >= 11
+#  pragma GCC diagnostic pop
+#endif
+}
+
+static inline void
+_Py_INCREF_CODE(PyCodeObject *co)
+{
+    _Py_THREAD_INCREF_OBJECT((PyObject *)co, co->_co_unique_id);
+}
 
+static inline void
+_Py_THREAD_DECREF_OBJECT(PyObject *obj, Py_ssize_t unique_id)
+{
     _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
-    PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
 
     // Unsigned comparison so that `unique_id=-1`, which indicates that
-    // per-thread refcounting has been disabled on this type, is handled by
+    // per-thread refcounting has been disabled on this object, is handled by
     // the "else".
-    if ((size_t)ht->unique_id < (size_t)tstate->refcounts.size) {
+    if ((size_t)unique_id < (size_t)tstate->refcounts.size) {
 #  ifdef Py_REF_DEBUG
-        _Py_INCREF_IncRefTotal();
+        _Py_DECREF_DecRefTotal();
 #  endif
-        _Py_INCREF_STAT_INC();
-        tstate->refcounts.values[ht->unique_id]++;
+        _Py_DECREF_STAT_INC();
+        tstate->refcounts.values[unique_id]--;
     }
     else {
-        // The slow path resizes the thread-local refcount array if necessary.
-        // It handles the unique_id=-1 case to keep the inlinable function smaller.
-        _PyType_IncrefSlow(ht);
+        // Directly decref the object if the id is not assigned or if
+        // per-thread refcounting has been disabled on this object.
+        Py_DECREF(obj);
     }
-
-#if defined(__GNUC__) && __GNUC__ >= 11
-#  pragma GCC diagnostic pop
-#endif
 }
 
 static inline void
@@ -341,25 +374,14 @@ _Py_DECREF_TYPE(PyTypeObject *type)
         _Py_DECREF_IMMORTAL_STAT_INC();
         return;
     }
-
-    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
     PyHeapTypeObject *ht = (PyHeapTypeObject *)type;
+    _Py_THREAD_DECREF_OBJECT((PyObject *)type, ht->unique_id);
+}
 
-    // Unsigned comparison so that `unique_id=-1`, which indicates that
-    // per-thread refcounting has been disabled on this type, is handled by
-    // the "else".
-    if ((size_t)ht->unique_id < (size_t)tstate->refcounts.size) {
-#  ifdef Py_REF_DEBUG
-        _Py_DECREF_DecRefTotal();
-#  endif
-        _Py_DECREF_STAT_INC();
-        tstate->refcounts.values[ht->unique_id]--;
-    }
-    else {
-        // Directly decref the type if the type id is not assigned or if
-        // per-thread refcounting has been disabled on this type.
-        Py_DECREF(type);
-    }
+static inline void
+_Py_DECREF_CODE(PyCodeObject *co)
+{
+    _Py_THREAD_DECREF_OBJECT((PyObject *)co, co->_co_unique_id);
 }
 #endif
 
diff --git a/Include/internal/pycore_uniqueid.h b/Include/internal/pycore_uniqueid.h
index 8f3b441..ad5dd38 100644
--- a/Include/internal/pycore_uniqueid.h
+++ b/Include/internal/pycore_uniqueid.h
@@ -49,7 +49,7 @@ struct _Py_unique_id_pool {
 extern Py_ssize_t _PyObject_AssignUniqueId(PyObject *obj);
 
 // Releases the allocated id back to the pool.
-extern void _PyObject_ReleaseUniqueId(Py_ssize_t unique_id);
+extern void _PyObject_DisablePerThreadRefcounting(PyObject *obj);
 
 // Merges the per-thread reference counts into the corresponding objects.
 extern void _PyObject_MergePerThreadRefcounts(_PyThreadStateImpl *tstate);
@@ -61,8 +61,8 @@ extern void _PyObject_FinalizePerThreadRefcounts(_PyThreadStateImpl *tstate);
 // Frees the interpreter's pool of type ids.
 extern void _PyObject_FinalizeUniqueIdPool(PyInterpreterState *interp);
 
-// Increfs the type, resizing the per-thread refcount array if necessary.
-PyAPI_FUNC(void) _PyType_IncrefSlow(PyHeapTypeObject *type);
+// Increfs the object, resizing the thread-local refcount array if necessary.
+PyAPI_FUNC(void) _PyObject_ThreadIncrefSlow(PyObject *obj, Py_ssize_t unique_id);
 
 #endif   /* Py_GIL_DISABLED */