gh-124218: Avoid refcount contention on builtins module (GH-125847)

colesbury · web-flow · commit 3c4a7fa6178d · 2024-10-24T12:44:38.000-04:00
This replaces `_PyEval_BuiltinsFromGlobals` with
`_PyDict_LoadBuiltinsFromGlobals`, which returns a new reference
instead of a borrowed reference. Internally, the new function uses
per-thread reference counting when possible to avoid contention on the
refcount fields on the builtins module.
diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h
@@ -83,9 +83,6 @@ extern void _PyEval_Fini(void);
 
 
 extern PyObject* _PyEval_GetBuiltins(PyThreadState *tstate);
-extern PyObject* _PyEval_BuiltinsFromGlobals(
-    PyThreadState *tstate,
-    PyObject *globals);
 
 // Trampoline API
 
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -108,6 +108,9 @@ extern Py_ssize_t _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyObject
 PyAPI_FUNC(PyObject *)_PyDict_LoadGlobal(PyDictObject *, PyDictObject *, PyObject *);
 PyAPI_FUNC(void) _PyDict_LoadGlobalStackRef(PyDictObject *, PyDictObject *, PyObject *, _PyStackRef *);
 
+// Loads the __builtins__ object from the globals dict. Returns a new reference.
+extern PyObject *_PyDict_LoadBuiltinsFromGlobals(PyObject *globals);
+
 /* Consumes references to key and value */
 PyAPI_FUNC(int) _PyDict_SetItem_Take2(PyDictObject *op, PyObject *key, PyObject *value);
 extern int _PyDict_SetItem_LockHeld(PyDictObject *dict, PyObject *name, PyObject *value);
@@ -318,6 +321,8 @@ PyDictObject *_PyObject_MaterializeManagedDict_LockHeld(PyObject *);
 #ifndef Py_GIL_DISABLED
 #  define _Py_INCREF_DICT Py_INCREF
 #  define _Py_DECREF_DICT Py_DECREF
+#  define _Py_INCREF_BUILTINS Py_INCREF
+#  define _Py_DECREF_BUILTINS Py_DECREF
 #else
 static inline Py_ssize_t
 _PyDict_UniqueId(PyDictObject *mp)
@@ -341,6 +346,30 @@ _Py_DECREF_DICT(PyObject *op)
     Py_ssize_t id = _PyDict_UniqueId((PyDictObject *)op);
     _Py_THREAD_DECREF_OBJECT(op, id);
 }
+
+// Like `_Py_INCREF_DICT`, but also handles non-dict objects because builtins
+// may not be a dict.
+static inline void
+_Py_INCREF_BUILTINS(PyObject *op)
+{
+    if (PyDict_CheckExact(op)) {
+        _Py_INCREF_DICT(op);
+    }
+    else {
+        Py_INCREF(op);
+    }
+}
+
+static inline void
+_Py_DECREF_BUILTINS(PyObject *op)
+{
+    if (PyDict_CheckExact(op)) {
+        _Py_DECREF_DICT(op);
+    }
+    else {
+        Py_DECREF(op);
+    }
+}
 #endif
 
 #ifdef __cplusplus
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -2511,6 +2511,40 @@ _PyDict_LoadGlobalStackRef(PyDictObject *globals, PyDictObject *builtins, PyObje
     assert(ix >= 0 || PyStackRef_IsNull(*res));
 }
 
+PyObject *
+_PyDict_LoadBuiltinsFromGlobals(PyObject *globals)
+{
+    if (!PyDict_Check(globals)) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    PyDictObject *mp = (PyDictObject *)globals;
+    PyObject *key = &_Py_ID(__builtins__);
+    Py_hash_t hash = unicode_get_hash(key);
+
+    // Use the stackref variant to avoid reference count contention on the
+    // builtins module in the free threading build. It's important not to
+    // make any escaping calls between the lookup and the `PyStackRef_CLOSE()`
+    // because the `ref` is not visible to the GC.
+    _PyStackRef ref;
+    Py_ssize_t ix = _Py_dict_lookup_threadsafe_stackref(mp, key, hash, &ref);
+    if (ix == DKIX_ERROR) {
+        return NULL;
+    }
+    if (PyStackRef_IsNull(ref)) {
+        return Py_NewRef(PyEval_GetBuiltins());
+    }
+    PyObject *builtins = PyStackRef_AsPyObjectBorrow(ref);
+    if (PyModule_Check(builtins)) {
+        builtins = _PyModule_GetDict(builtins);
+        assert(builtins != NULL);
+    }
+    _Py_INCREF_BUILTINS(builtins);
+    PyStackRef_CLOSE(ref);
+    return builtins;
+}
+
 /* Consumes references to key and value */
 static int
 setitem_take2_lock_held(PyDictObject *mp, PyObject *key, PyObject *value)
diff --git a/Objects/frameobject.c b/Objects/frameobject.c
@@ -1,8 +1,9 @@
 /* Frame object implementation */
 
 #include "Python.h"
-#include "pycore_ceval.h"         // _PyEval_BuiltinsFromGlobals()
+#include "pycore_ceval.h"         // _PyEval_SetOpcodeTrace()
 #include "pycore_code.h"          // CO_FAST_LOCAL, etc.
+#include "pycore_dict.h"          // _PyDict_LoadBuiltinsFromGlobals()
 #include "pycore_function.h"      // _PyFunction_FromConstructor()
 #include "pycore_moduleobject.h"  // _PyModule_GetDict()
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
@@ -1899,7 +1900,7 @@ PyFrameObject*
 PyFrame_New(PyThreadState *tstate, PyCodeObject *code,
             PyObject *globals, PyObject *locals)
 {
-    PyObject *builtins = _PyEval_BuiltinsFromGlobals(tstate, globals); // borrowed ref
+    PyObject *builtins = _PyDict_LoadBuiltinsFromGlobals(globals);
     if (builtins == NULL) {
         return NULL;
     }
@@ -1914,6 +1915,7 @@ PyFrame_New(PyThreadState *tstate, PyCodeObject *code,
         .fc_closure = NULL
     };
     PyFunctionObject *func = _PyFunction_FromConstructor(&desc);
+    _Py_DECREF_BUILTINS(builtins);
     if (func == NULL) {
         return NULL;
     }
@@ -2204,21 +2206,3 @@ PyFrame_GetGenerator(PyFrameObject *frame)
     PyGenObject *gen = _PyGen_GetGeneratorFromFrame(frame->f_frame);
     return Py_NewRef(gen);
 }
-
-PyObject*
-_PyEval_BuiltinsFromGlobals(PyThreadState *tstate, PyObject *globals)
-{
-    PyObject *builtins = PyDict_GetItemWithError(globals, &_Py_ID(__builtins__));
-    if (builtins) {
-        if (PyModule_Check(builtins)) {
-            builtins = _PyModule_GetDict(builtins);
-            assert(builtins != NULL);
-        }
-        return builtins;
-    }
-    if (PyErr_Occurred()) {
-        return NULL;
-    }
-
-    return _PyEval_GetBuiltins(tstate);
-}
diff --git a/Objects/funcobject.c b/Objects/funcobject.c
@@ -2,7 +2,6 @@
 /* Function object implementation */
 
 #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()
@@ -115,12 +114,7 @@ _PyFunction_FromConstructor(PyFrameConstructor *constr)
     }
     _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);
-    }
+    _Py_INCREF_BUILTINS(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);
@@ -153,8 +147,6 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
     assert(PyDict_Check(globals));
     _Py_INCREF_DICT(globals);
 
-    PyThreadState *tstate = _PyThreadState_GET();
-
     PyCodeObject *code_obj = (PyCodeObject *)code;
     _Py_INCREF_CODE(code_obj);
 
@@ -188,16 +180,10 @@ PyFunction_NewWithQualName(PyObject *code, PyObject *globals, PyObject *qualname
         goto error;
     }
 
-    builtins = _PyEval_BuiltinsFromGlobals(tstate, globals); // borrowed ref
+    builtins = _PyDict_LoadBuiltinsFromGlobals(globals);
     if (builtins == NULL) {
         goto error;
     }
-    if (PyDict_Check(builtins)) {
-        _Py_INCREF_DICT(builtins);
-    }
-    else {
-        Py_INCREF(builtins);
-    }
 
     PyFunctionObject *op = PyObject_GC_New(PyFunctionObject, &PyFunction_Type);
     if (op == NULL) {
@@ -1078,12 +1064,7 @@ func_clear(PyObject *self)
     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_DECREF_BUILTINS(builtins);
     }
     Py_CLEAR(op->func_module);
     Py_CLEAR(op->func_defaults);
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -639,7 +639,7 @@ PyEval_EvalCode(PyObject *co, PyObject *globals, PyObject *locals)
     if (locals == NULL) {
         locals = globals;
     }
-    PyObject *builtins = _PyEval_BuiltinsFromGlobals(tstate, globals); // borrowed ref
+    PyObject *builtins = _PyDict_LoadBuiltinsFromGlobals(globals);
     if (builtins == NULL) {
         return NULL;
     }
@@ -654,6 +654,7 @@ PyEval_EvalCode(PyObject *co, PyObject *globals, PyObject *locals)
         .fc_closure = NULL
     };
     PyFunctionObject *func = _PyFunction_FromConstructor(&desc);
+    _Py_DECREF_BUILTINS(builtins);
     if (func == NULL) {
         return NULL;
     }
@@ -1899,7 +1900,7 @@ PyEval_EvalCodeEx(PyObject *_co, PyObject *globals, PyObject *locals,
     if (defaults == NULL) {
         return NULL;
     }
-    PyObject *builtins = _PyEval_BuiltinsFromGlobals(tstate, globals); // borrowed ref
+    PyObject *builtins = _PyDict_LoadBuiltinsFromGlobals(globals);
     if (builtins == NULL) {
         Py_DECREF(defaults);
         return NULL;
@@ -1954,6 +1955,7 @@ PyEval_EvalCodeEx(PyObject *_co, PyObject *globals, PyObject *locals,
     Py_XDECREF(func);
     Py_XDECREF(kwnames);
     PyMem_Free(newargs);
+    _Py_DECREF_BUILTINS(builtins);
     Py_DECREF(defaults);
     return res;
 }
