Lock dict when specializing module lookup

DinoV · DinoV · commit a857147da953 · 2024-01-30T16:07:51.000-08:00
diff --git a/Python/specialize.c b/Python/specialize.c
@@ -3,6 +3,7 @@
 #include "opcode.h"
 
 #include "pycore_code.h"
+#include "pycore_critical_section.h"    // Py_BEGIN_CRITICAL_SECTION, Py_END_CRITICAL_SECTION
 #include "pycore_descrobject.h"   // _PyMethodWrapper_Type
 #include "pycore_dict.h"          // DICT_KEYS_UNICODE
 #include "pycore_function.h"      // _PyFunction_GetVersionForCurrentState()
@@ -600,17 +601,11 @@ static uint32_t function_get_version(PyObject *o, int opcode);
 static uint32_t type_get_version(PyTypeObject *t, int opcode);
 
 static int
-specialize_module_load_attr(
-    PyObject *owner, _Py_CODEUNIT *instr, PyObject *name
-) {
+specialize_module_load_attr_lock_held(
+    PyDictObject *dict, _Py_CODEUNIT *instr, PyObject *name
+)
+{
     _PyAttrCache *cache = (_PyAttrCache *)(instr + 1);
-    PyModuleObject *m = (PyModuleObject *)owner;
-    assert((owner->ob_type->tp_flags & Py_TPFLAGS_MANAGED_DICT) == 0);
-    PyDictObject *dict = (PyDictObject *)m->md_dict;
-    if (dict == NULL) {
-        SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_NO_DICT);
-        return -1;
-    }
     if (dict->ma_keys->dk_kind != DICT_KEYS_UNICODE) {
         SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_ATTR_NON_STRING_OR_SPLIT);
         return -1;
@@ -642,6 +637,23 @@ specialize_module_load_attr(
     return 0;
 }
 
+static int
+specialize_module_load_attr(
+    PyObject *owner, _Py_CODEUNIT *instr, PyObject *name
+) {
+    PyModuleObject *m = (PyModuleObject *)owner;
+    assert((owner->ob_type->tp_flags & Py_TPFLAGS_MANAGED_DICT) == 0);
+    PyDictObject *dict = (PyDictObject *)m->md_dict;
+    if (dict == NULL) {
+        SPECIALIZATION_FAIL(LOAD_ATTR, SPEC_FAIL_NO_DICT);
+        return -1;
+    }
+    int res;
+    Py_BEGIN_CRITICAL_SECTION(dict);
+    res = specialize_module_load_attr_lock_held(dict, instr, name);
+    Py_END_CRITICAL_SECTION();
+    return res;
+}
 
 
 /* Attribute specialization */
@@ -836,18 +848,34 @@ specialize_dict_access(
             SPECIALIZATION_FAIL(base_op, SPEC_FAIL_ATTR_NON_STRING_OR_SPLIT);
             return 0;
         }
+
+#ifdef Py_GIL_DISABLED
+        unsigned int version = _Py_atomic_load_uint(&type->tp_version_tag);
+#else
+        unsigned int version = dict->ma_version_tag;
+#endif
+
+        int res;
+        Py_BEGIN_CRITICAL_SECTION(dict);
+
         Py_ssize_t index =
             _PyDict_LookupIndex(dict, name);
         if (index != (uint16_t)index) {
             SPECIALIZATION_FAIL(base_op,
                                 index == DKIX_EMPTY ?
                                 SPEC_FAIL_ATTR_NOT_IN_DICT :
                                 SPEC_FAIL_OUT_OF_RANGE);
-            return 0;
+            res = 0;
+            goto exit;
         }
         cache->index = (uint16_t)index;
-        write_u32(cache->version, type->tp_version_tag);
+        write_u32(cache->version, version);
         instr->op.code = hint_op;
+
+        res = 1;
+exit:
+        Py_END_CRITICAL_SECTION();
+        return res;
     }
     return 1;
 }
