Make _PyDict_LoadGlobal threadsafe

DinoV · DinoV · commit 548a7c275ccf · 2024-04-03T17:34:50.000-07:00
diff --git 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
@@ -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;
@@ -2343,11 +2343,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;
@@ -2358,17 +2359,31 @@ _PyDict_LoadGlobal(PyDictObject *globals, PyDictObject *builtins, PyObject *key)
             return NULL;
     }
 
+#ifdef Py_GIL_DISABLED
     /* 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;
+#else
+    /* namespace 1: globals */
+    ix = _Py_dict_lookup(globals, key, hash, &value);
+    if (ix == DKIX_ERROR)
+        return NULL;
+    if (ix != DKIX_EMPTY && value != NULL)
+        return Py_NewRef(value);
+
+    /* namespace 2: builtins */
+    ix = _Py_dict_lookup(builtins, key, hash, &value);
+    assert(ix >= 0 || value == NULL);
+    return Py_XNewRef(value);
+#endif
 }
 
 /* Consumes references to key and value */
diff --git 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
@@ -1426,7 +1426,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
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h

Original file line number	Diff line number	Diff line change
`@@ -1426,7 +1426,6 @@ dummy_func(`
`1426`	`1426`	`}`
`1427`	`1427`	`ERROR_IF(true, error);`
`1428`	`1428`	`}`
`1429`		`- Py_INCREF(res);`
`1430`	`1429`	`}`
`1431`	`1430`	`else {`
`1432`	`1431`	`/* Slow-path if globals or builtins is not a dict */`