gh-112075: Make _PyDict_LoadGlobal thread safe (#117529)

Make _PyDict_LoadGlobal threadsafe

6 files changed, 23 insertions, 29 deletions

diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
index 5507bdd..fba0dfc 100644
--- a/Include/internal/pycore_dict.h
+++ b/Include/internal/pycore_dict.h
@@ -97,6 +97,7 @@ extern void _PyDictKeys_DecRef(PyDictKeysObject *keys);
  * -1 when no entry found, -3 when compare raises error.
  */
 extern Py_ssize_t _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
+extern Py_ssize_t _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr);
 
 extern Py_ssize_t _PyDict_LookupIndex(PyDictObject *, PyObject *);
 extern Py_ssize_t _PyDictKeys_StringLookup(PyDictKeysObject* dictkeys, PyObject *key);
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index 58a3d97..b62d39a 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -1065,7 +1065,6 @@ compare_unicode_generic(PyDictObject *mp, PyDictKeysObject *dk,
     assert(ep->me_key != NULL);
     assert(PyUnicode_CheckExact(ep->me_key));
     assert(!PyUnicode_CheckExact(key));
-    // TODO: Thread safety
 
     if (unicode_get_hash(ep->me_key) == hash) {
         PyObject *startkey = ep->me_key;
@@ -1192,7 +1191,8 @@ _Py_dict_lookup(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **valu
     PyDictKeysObject *dk;
     DictKeysKind kind;
     Py_ssize_t ix;
-    // TODO: Thread safety
+
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(mp);
 start:
     dk = mp->ma_keys;
     kind = dk->dk_kind;
@@ -1390,7 +1390,7 @@ dictkeys_generic_lookup_threadsafe(PyDictObject *mp, PyDictKeysObject* dk, PyObj
     return do_lookup(mp, dk, key, hash, compare_generic_threadsafe);
 }
 
-static Py_ssize_t
+Py_ssize_t
 _Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr)
 {
     PyDictKeysObject *dk;
@@ -1488,6 +1488,16 @@ read_failed:
     return ix;
 }
 
+#else   // Py_GIL_DISABLED
+
+Py_ssize_t
+_Py_dict_lookup_threadsafe(PyDictObject *mp, PyObject *key, Py_hash_t hash, PyObject **value_addr)
+{
+    Py_ssize_t ix = _Py_dict_lookup(mp, key, hash, value_addr);
+    Py_XNewRef(*value_addr);
+    return ix;
+}
+
 #endif
 
 int
@@ -2343,11 +2353,12 @@ _PyDict_GetItemStringWithError(PyObject *v, const char *key)
  * Raise an exception and return NULL if an error occurred (ex: computing the
  * key hash failed, key comparison failed, ...). Return NULL if the key doesn't
  * exist. Return the value if the key exists.
+ *
+ * Returns a new reference.
  */
 PyObject *
 _PyDict_LoadGlobal(PyDictObject *globals, PyDictObject *builtins, PyObject *key)
 {
-    // TODO: Thread safety
     Py_ssize_t ix;
     Py_hash_t hash;
     PyObject *value;
@@ -2359,14 +2370,14 @@ _PyDict_LoadGlobal(PyDictObject *globals, PyDictObject *builtins, PyObject *key)
     }
 
     /* namespace 1: globals */
-    ix = _Py_dict_lookup(globals, key, hash, &value);
+    ix = _Py_dict_lookup_threadsafe(globals, key, hash, &value);
     if (ix == DKIX_ERROR)
         return NULL;
     if (ix != DKIX_EMPTY && value != NULL)
         return value;
 
     /* namespace 2: builtins */
-    ix = _Py_dict_lookup(builtins, key, hash, &value);
+    ix = _Py_dict_lookup_threadsafe(builtins, key, hash, &value);
     assert(ix >= 0 || value == NULL);
     return value;
 }
@@ -3214,11 +3225,7 @@ dict_subscript(PyObject *self, PyObject *key)
         if (hash == -1)
             return NULL;
     }
-#ifdef Py_GIL_DISABLED
     ix = _Py_dict_lookup_threadsafe(mp, key, hash, &value);
-#else
-    ix = _Py_dict_lookup(mp, key, hash, &value);
-#endif
     if (ix == DKIX_ERROR)
         return NULL;
     if (ix == DKIX_EMPTY || value == NULL) {
@@ -3238,11 +3245,7 @@ dict_subscript(PyObject *self, PyObject *key)
         _PyErr_SetKeyError(key);
         return NULL;
     }
-#ifdef Py_GIL_DISABLED
     return value;
-#else
-    return Py_NewRef(value);
-#endif
 }
 
 static int
@@ -4109,24 +4112,13 @@ dict_get_impl(PyDictObject *self, PyObject *key, PyObject *default_value)
         if (hash == -1)
             return NULL;
     }
-#ifdef Py_GIL_DISABLED
     ix = _Py_dict_lookup_threadsafe(self, key, hash, &val);
-#else
-    ix = _Py_dict_lookup(self, key, hash, &val);
-#endif
     if (ix == DKIX_ERROR)
         return NULL;
-#ifdef Py_GIL_DISABLED
     if (ix == DKIX_EMPTY || val == NULL) {
         val = Py_NewRef(default_value);
     }
     return val;
-#else
-    if (ix == DKIX_EMPTY || val == NULL) {
-        val = default_value;
-    }
-    return Py_NewRef(val);
-#endif
 }
 
 static int
diff --git a/Objects/odictobject.c b/Objects/odictobject.c
index 421bc52..53f64fc 100644
--- a/Objects/odictobject.c
+++ b/Objects/odictobject.c
@@ -535,8 +535,12 @@ _odict_get_index_raw(PyODictObject *od, PyObject *key, Py_hash_t hash)
     PyObject *value = NULL;
     PyDictKeysObject *keys = ((PyDictObject *)od)->ma_keys;
     Py_ssize_t ix;
-
+#ifdef Py_GIL_DISABLED
+    ix = _Py_dict_lookup_threadsafe((PyDictObject *)od, key, hash, &value);
+    Py_XDECREF(value);
+#else
     ix = _Py_dict_lookup((PyDictObject *)od, key, hash, &value);
+#endif
     if (ix == DKIX_EMPTY) {
         return keys->dk_nentries;  /* index of new entry */
     }
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index 8af48d9..d6fb66a 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -1427,7 +1427,6 @@ dummy_func(
                     }
                     ERROR_IF(true, error);
                 }
-                Py_INCREF(res);
             }
             else {
                 /* Slow-path if globals or builtins is not a dict */
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index 8c3d41b..a6e28f6 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -1254,7 +1254,6 @@
                     }
                     if (true) JUMP_TO_ERROR();
                 }
-                Py_INCREF(res);
             }
             else {
                 /* Slow-path if globals or builtins is not a dict */
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index 0116acd..a7764b0 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -4256,7 +4256,6 @@
                         }
                         if (true) goto error;
                     }
-                    Py_INCREF(res);
                 }
                 else {
                     /* Slow-path if globals or builtins is not a dict */