WIP thread refcount

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);

Objects/codeobject.c

@@ -1860,9 +1860,9 @@ code_dealloc(PyCodeObject *co)
     Py_XDECREF(co->co_qualname);
     Py_XDECREF(co->co_linetable);
     Py_XDECREF(co->co_exceptiontable);
-    if (co->_co_unique_id != -1) {
-        _PyObject_ReleaseUniqueId(co->_co_unique_id);
-    }
+#ifdef Py_GIL_DISABLED
+    assert(co->_co_unique_id == -1);
+#endif
     if (co->_co_cached != NULL) {
         Py_XDECREF(co->_co_cached->_co_code);
         Py_XDECREF(co->_co_cached->_co_cellvars);

Objects/typeobject.c

@@ -6080,7 +6080,7 @@ type_dealloc(PyObject *self)
     Py_XDECREF(et->ht_module);
     PyMem_Free(et->_ht_tpname);
 #ifdef Py_GIL_DISABLED
-    _PyObject_ReleaseUniqueId(et->unique_id);
+    assert(et->unique_id == -1);
 #endif
     et->ht_token = NULL;
     Py_TYPE(type)->tp_free((PyObject *)type);

Python/gc_free_threading.c

@@ -215,19 +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 and code objects also use per-thread refcounting, which
-    // should also be disabled when we turn off deferred refcounting.
-    if (PyType_Check(op) && PyType_HasFeature((PyTypeObject *)op, Py_TPFLAGS_HEAPTYPE)) {
-        PyHeapTypeObject *ht = (PyHeapTypeObject *)op;
-        _PyObject_ReleaseUniqueId(ht->unique_id);
-        ht->unique_id = -1;
-    }
-    else if (PyCode_Check(op)) {
-        PyCodeObject *co = (PyCodeObject *)op;
-        _PyObject_ReleaseUniqueId(co->_co_unique_id);
-        co->_co_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/uniqueid.c

@@ -98,24 +98,48 @@ _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 *)op;
+            id = ht->unique_id;
+            ht->unique_id = -1;
+        }
+    }
+    else if (PyCode_Check(obj)) {
+        PyCodeObject *co = (PyCodeObject *)op;
+        id = co->unique_id;
+        co->_co_unique_id = -1;
+    }
+    return id;
+}
+
+void
+_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)
 {
@@ -179,24 +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 case 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 if (PyCode_Check(obj)) {
-            PyCodeObject *co = (PyCodeObject *)obj;
-            co->_co_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);