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.

Python/gc_free_threading.c

@@ -15,7 +15,7 @@
 #include "pycore_tstate.h"        // _PyThreadStateImpl
 #include "pycore_weakref.h"       // _PyWeakref_ClearRef()
 #include "pydtrace.h"
-#include "pycore_uniqueid.h"      // _PyType_MergeThreadLocalRefcounts
+#include "pycore_uniqueid.h"      // _PyObject_MergeThreadLocalRefcounts()
 
 #ifdef Py_GIL_DISABLED
 
@@ -215,15 +215,10 @@ disable_deferred_refcounting(PyObject *op)
         op->ob_gc_bits &= ~_PyGC_BITS_DEFERRED;
         op->ob_ref_shared -= _Py_REF_SHARED(_Py_REF_DEFERRED, 0);
         merge_refcount(op, 0);
-    }
 
-    // Heap types also use per-thread refcounting -- disable it here.
-    if (PyType_Check(op)) {
-        if (PyType_HasFeature((PyTypeObject *)op, Py_TPFLAGS_HEAPTYPE)) {
-            PyHeapTypeObject *ht = (PyHeapTypeObject *)op;
-            _PyObject_ReleaseUniqueId(ht->unique_id);
-            ht->unique_id = -1;
-        }
+        // Heap types and code objects also use per-thread refcounting, which
+        // should also be disabled when we turn off deferred refcounting.
+        _PyObject_DisablePerThreadRefcounting(op);
     }
 
     // Generators and frame objects may contain deferred references to other

Python/pylifecycle.c

@@ -28,7 +28,7 @@
 #include "pycore_sliceobject.h"   // _PySlice_Fini()
 #include "pycore_sysmodule.h"     // _PySys_ClearAuditHooks()
 #include "pycore_traceback.h"     // _Py_DumpTracebackThreads()
-#include "pycore_uniqueid.h"      // _PyType_FinalizeIdPool()
+#include "pycore_uniqueid.h"      // _PyObject_FinalizeUniqueIdPool()
 #include "pycore_typeobject.h"    // _PyTypes_InitTypes()
 #include "pycore_typevarobject.h" // _Py_clear_generic_types()
 #include "pycore_unicodeobject.h" // _PyUnicode_InitTypes()

Python/pystate.c

@@ -20,7 +20,7 @@
 #include "pycore_runtime_init.h"  // _PyRuntimeState_INIT
 #include "pycore_sysmodule.h"     // _PySys_Audit()
 #include "pycore_obmalloc.h"      // _PyMem_obmalloc_state_on_heap()
-#include "pycore_uniqueid.h"      // _PyType_FinalizeThreadLocalRefcounts()
+#include "pycore_uniqueid.h"      // _PyObject_FinalizePerThreadRefcounts()
 
 /* --------------------------------------------------------------------------
 CAUTION

Python/uniqueid.c

@@ -98,36 +98,60 @@ _PyObject_AssignUniqueId(PyObject *obj)
     return unique_id;
 }
 
-void
-_PyObject_ReleaseUniqueId(Py_ssize_t unique_id)
+static void
+release_unique_id(Py_ssize_t unique_id)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
     struct _Py_unique_id_pool *pool = &interp->unique_ids;
 
-    if (unique_id < 0) {
-        // The id is not assigned
-        return;
-    }
-
     LOCK_POOL(pool);
+    assert(unique_id >= 0 && unique_id < pool->size);
     _Py_unique_id_entry *entry = &pool->table[unique_id];
     entry->next = pool->freelist;
     pool->freelist = entry;
     UNLOCK_POOL(pool);
 }
 
+static Py_ssize_t
+clear_unique_id(PyObject *obj)
+{
+    Py_ssize_t id = -1;
+    if (PyType_Check(obj)) {
+        if (PyType_HasFeature((PyTypeObject *)obj, Py_TPFLAGS_HEAPTYPE)) {
+            PyHeapTypeObject *ht = (PyHeapTypeObject *)obj;
+            id = ht->unique_id;
+            ht->unique_id = -1;
+        }
+    }
+    else if (PyCode_Check(obj)) {
+        PyCodeObject *co = (PyCodeObject *)obj;
+        id = co->_co_unique_id;
+        co->_co_unique_id = -1;
+    }
+    return id;
+}
+
 void
-_PyType_IncrefSlow(PyHeapTypeObject *type)
+_PyObject_DisablePerThreadRefcounting(PyObject *obj)
+{
+    Py_ssize_t id = clear_unique_id(obj);
+    if (id >= 0) {
+        release_unique_id(id);
+    }
+}
+
+void
+_PyObject_ThreadIncrefSlow(PyObject *obj, Py_ssize_t unique_id)
 {
     _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
-    if (type->unique_id < 0 || resize_local_refcounts(tstate) < 0) {
-        // just incref the type directly.
-        Py_INCREF(type);
+    if (unique_id < 0 || resize_local_refcounts(tstate) < 0) {
+        // just incref the object directly.
+        Py_INCREF(obj);
         return;
     }
 
-    assert(type->unique_id < tstate->refcounts.size);
-    tstate->refcounts.values[type->unique_id]++;
+    assert(unique_id < tstate->refcounts.size);
+    tstate->refcounts.values[unique_id]++;
 #ifdef Py_REF_DEBUG
     _Py_IncRefTotal((PyThreadState *)tstate);
 #endif
@@ -179,20 +203,15 @@ _PyObject_FinalizeUniqueIdPool(PyInterpreterState *interp)
         pool->freelist = next;
     }
 
-    // Now everything non-NULL is a type. Set the type's id to -1 in case it
-    // outlives the interpreter.
+    // Now everything non-NULL is a object. Clear their unique ids as the
+    // object outlives the interpreter.
     for (Py_ssize_t i = 0; i < pool->size; i++) {
         PyObject *obj = pool->table[i].obj;
         pool->table[i].obj = NULL;
-        if (obj == NULL) {
-            continue;
-        }
-        if (PyType_Check(obj)) {
-            assert(PyType_HasFeature((PyTypeObject *)obj, Py_TPFLAGS_HEAPTYPE));
-            ((PyHeapTypeObject *)obj)->unique_id = -1;
-        }
-        else {
-            Py_UNREACHABLE();
+        if (obj != NULL) {
+            Py_ssize_t id = clear_unique_id(obj);
+            (void)id;
+            assert(id == i);
         }
     }
     PyMem_Free(pool->table);