pythongh-124218: Use per-thread reference counting for globals and bu…

…iltins

Use per-thread refcounting for the reference from function objects to
the globals and builtins dictionaries.

Include/internal/pycore_dict.h

@@ -229,6 +229,8 @@ static inline PyDictUnicodeEntry* DK_UNICODE_ENTRIES(PyDictKeysObject *dk) {
 #define DICT_VERSION_INCREMENT (1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS))
 #define DICT_WATCHER_MASK ((1 << DICT_MAX_WATCHERS) - 1)
 #define DICT_WATCHER_AND_MODIFICATION_MASK ((1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS)) - 1)
+#define DICT_UNIQUE_ID_SHIFT (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS)
+#define DICT_UNIQUE_ID_MAX (UINT64_MAX >> DICT_UNIQUE_ID_SHIFT)
 
 
 PyAPI_FUNC(void)
@@ -307,8 +309,39 @@ _PyInlineValuesSize(PyTypeObject *tp)
 int
 _PyDict_DetachFromObject(PyDictObject *dict, PyObject *obj);
 
+// Enables per-thread ref counting on this dict in the free threading build
+extern void _PyDict_EnablePerThreadRefcounting(PyObject *op);
+
 PyDictObject *_PyObject_MaterializeManagedDict_LockHeld(PyObject *);
 
+#ifndef Py_GIL_DISABLED
+#  define _Py_INCREF_DICT Py_INCREF
+#  define _Py_DECREF_DICT Py_DECREF
+#else
+static inline Py_ssize_t
+_PyDict_UniqueId(PyDictObject *mp)
+{
+    // Offset by on so that _ma_watcher_tag=0 represents an unassigned id
+    return (Py_ssize_t)(mp->_ma_watcher_tag >> DICT_UNIQUE_ID_SHIFT) - 1;
+}
+
+static inline void
+_Py_INCREF_DICT(PyObject *op)
+{
+    assert(PyDict_Check(op));
+    Py_ssize_t id = _PyDict_UniqueId((PyDictObject *)op);
+    _Py_THREAD_INCREF_OBJECT(op, id);
+}
+
+static inline void
+_Py_DECREF_DICT(PyObject *op)
+{
+    assert(PyDict_Check(op));
+    Py_ssize_t id = _PyDict_UniqueId((PyDictObject *)op);
+    _Py_THREAD_DECREF_OBJECT(op, id);
+}
+#endif
+
 #ifdef __cplusplus
 }
 #endif

Objects/dictobject.c

@@ -1636,6 +1636,20 @@ _PyDict_MaybeUntrack(PyObject *op)
     _PyObject_GC_UNTRACK(op);
 }
 
+void
+_PyDict_EnablePerThreadRefcounting(PyObject *op)
+{
+    assert(PyDict_Check(op));
+#ifdef Py_GIL_DISABLED
+    Py_ssize_t id = _PyObject_AssignUniqueId(op);
+
+    PyDictObject *mp = (PyDictObject *)op;
+    assert((mp->_ma_watcher_tag >> DICT_UNIQUE_ID_SHIFT) == 0);
+    // Plus 1 so that _ma_watcher_tag=0 represents an unassigned id
+    mp->_ma_watcher_tag += (id + 1) << DICT_UNIQUE_ID_SHIFT;
+#endif
+}
+
 static inline int
 is_unusable_slot(Py_ssize_t ix)
 {

Objects/funcobject.c

@@ -3,6 +3,7 @@
 
 #include "Python.h"
 #include "pycore_ceval.h"         // _PyEval_BuiltinsFromGlobals()
+#include "pycore_dict.h"          // _Py_INCREF_DICT()
 #include "pycore_long.h"          // _PyLong_GetOne()
 #include "pycore_modsupport.h"    // _PyArg_NoKeywords()
 #include "pycore_object.h"        // _PyObject_GC_UNTRACK()
@@ -112,8 +113,15 @@ _PyFunction_FromConstructor(PyFrameConstructor *constr)
         Py_XDECREF(module);
         return NULL;
     }
-    op->func_globals = Py_NewRef(constr->fc_globals);
-    op->func_builtins = Py_NewRef(constr->fc_builtins);
+    _Py_INCREF_DICT(constr->fc_globals);
+    op->func_globals = constr->fc_globals;
+    if (PyDict_Check(constr->fc_builtins)) {
+        _Py_INCREF_DICT(constr->fc_builtins);
+    }
+    else {
+        Py_INCREF(constr->fc_builtins);
+    }
+    op->func_builtins = constr->fc_builtins;
     op->func_name = Py_NewRef(constr->fc_name);
     op->func_qualname = Py_NewRef(constr->fc_qualname);
     _Py_INCREF_CODE((PyCodeObject *)constr->fc_code);
@@ -143,7 +151,7 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
 {
     assert(globals != NULL);
     assert(PyDict_Check(globals));
-    Py_INCREF(globals);
+    _Py_INCREF_DICT(globals);
 
     PyThreadState *tstate = _PyThreadState_GET();
 
@@ -184,7 +192,12 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     if (builtins == NULL) {
         goto error;
     }
-    Py_INCREF(builtins);
+    if (PyDict_Check(builtins)) {
+        _Py_INCREF_DICT(builtins);
+    }
+    else {
+        Py_INCREF(builtins);
+    }
 
     PyFunctionObject *op = PyObject_GC_New(PyFunctionObject, &PyFunction_Type);
     if (op == NULL) {
@@ -1057,8 +1070,21 @@ func_clear(PyObject *self)
 {
     PyFunctionObject *op = _PyFunction_CAST(self);
     func_clear_version(_PyInterpreterState_GET(), op);
-    Py_CLEAR(op->func_globals);
-    Py_CLEAR(op->func_builtins);
+    PyObject *globals = op->func_globals;
+    op->func_globals = NULL;
+    if (globals != NULL) {
+        _Py_DECREF_DICT(globals);
+    }
+    PyObject *builtins = op->func_builtins;
+    op->func_builtins = NULL;
+    if (builtins != NULL) {
+        if (PyDict_Check(builtins)) {
+            _Py_DECREF_DICT(builtins);
+        }
+        else {
+            Py_DECREF(builtins);
+        }
+    }
     Py_CLEAR(op->func_module);
     Py_CLEAR(op->func_defaults);
     Py_CLEAR(op->func_kwdefaults);

Objects/moduleobject.c

@@ -3,6 +3,7 @@
 
 #include "Python.h"
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
+#include "pycore_dict.h"          // _PyDict_EnablePerThreadRefcounting()
 #include "pycore_fileutils.h"     // _Py_wgetcwd
 #include "pycore_interp.h"        // PyInterpreterState.importlib
 #include "pycore_long.h"          // _PyLong_GetOne()
@@ -105,7 +106,7 @@ new_module_notrack(PyTypeObject *mt)
 static void
 track_module(PyModuleObject *m)
 {
-    _PyObject_SetDeferredRefcount(m->md_dict);
+    _PyDict_EnablePerThreadRefcounting(m->md_dict);
     PyObject_GC_Track(m->md_dict);
 
     _PyObject_SetDeferredRefcount((PyObject *)m);

Python/uniqueid.c

@@ -1,5 +1,6 @@
 #include "Python.h"
 
+#include "pycore_dict.h"        // DICT_UNIQUE_ID_MAX
 #include "pycore_lock.h"        // PyMutex_LockFlags()
 #include "pycore_pystate.h"     // _PyThreadState_GET()
 #include "pycore_object.h"      // _Py_IncRefTotal
@@ -90,10 +91,15 @@ _PyObject_AssignUniqueId(PyObject *obj)
     }
 
     _Py_unique_id_entry *entry = pool->freelist;
+    Py_ssize_t unique_id = (entry - pool->table);
+    if ((uint64_t)unique_id >= DICT_UNIQUE_ID_MAX) {
+        UNLOCK_POOL(pool);
+        return -1;
+    }
+
     pool->freelist = entry->next;
     entry->obj = obj;
     _PyObject_SetDeferredRefcount(obj);
-    Py_ssize_t unique_id = (entry - pool->table);
     UNLOCK_POOL(pool);
     return unique_id;
 }
@@ -128,6 +134,11 @@ clear_unique_id(PyObject *obj)
         id = co->_co_unique_id;
         co->_co_unique_id = -1;
     }
+    else if (PyDict_Check(obj)) {
+        PyDictObject *mp = (PyDictObject *)obj;
+        id = _PyDict_UniqueId(mp);
+        mp->_ma_watcher_tag &= ~(UINT64_MAX << DICT_UNIQUE_ID_SHIFT);
+    }
     return id;
 }