[gh-117657] Fix some issues with TSAN in typeobject (#118249)

DinoV · web-flow · commit 19d468a3ab1d · 2024-04-30T12:37:38.000-07:00
Fix some racing reads in typebobject.c
diff --git a/Include/internal/pycore_pyatomic_ft_wrappers.h b/Include/internal/pycore_pyatomic_ft_wrappers.h
@@ -43,6 +43,8 @@ extern "C" {
     _Py_atomic_load_uint8_relaxed(&value)
 #define FT_ATOMIC_LOAD_UINT16_RELAXED(value) \
     _Py_atomic_load_uint16_relaxed(&value)
+#define FT_ATOMIC_LOAD_UINT32_RELAXED(value) \
+    _Py_atomic_load_uint32_relaxed(&value)
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) \
     _Py_atomic_store_ptr_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) \
@@ -55,6 +57,9 @@ extern "C" {
     _Py_atomic_store_uint8_relaxed(&value, new_value)
 #define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) \
     _Py_atomic_store_uint16_relaxed(&value, new_value)
+#define FT_ATOMIC_STORE_UINT32_RELAXED(value, new_value) \
+    _Py_atomic_store_uint32_relaxed(&value, new_value)
+
 #else
 #define FT_ATOMIC_LOAD_PTR(value) value
 #define FT_ATOMIC_STORE_PTR(value, new_value) value = new_value
@@ -69,12 +74,15 @@ extern "C" {
 #define FT_ATOMIC_STORE_UINT8(value, new_value) value = new_value
 #define FT_ATOMIC_LOAD_UINT8_RELAXED(value) value
 #define FT_ATOMIC_LOAD_UINT16_RELAXED(value) value
+#define FT_ATOMIC_LOAD_UINT32_RELAXED(value) value
 #define FT_ATOMIC_STORE_PTR_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_PTR_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINTPTR_RELEASE(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_SSIZE_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINT8_RELAXED(value, new_value) value = new_value
 #define FT_ATOMIC_STORE_UINT16_RELAXED(value, new_value) value = new_value
+#define FT_ATOMIC_STORE_UINT32_RELAXED(value, new_value) value = new_value
+
 #endif
 
 #ifdef __cplusplus
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
@@ -43,7 +43,8 @@ class object "PyObject *" "&PyBaseObject_Type"
          & ((1 << MCACHE_SIZE_EXP) - 1))
 
 #define MCACHE_HASH_METHOD(type, name)                                  \
-    MCACHE_HASH((type)->tp_version_tag, ((Py_ssize_t)(name)) >> 3)
+    MCACHE_HASH(FT_ATOMIC_LOAD_UINT32_RELAXED((type)->tp_version_tag),   \
+                ((Py_ssize_t)(name)) >> 3)
 #define MCACHE_CACHEABLE_NAME(name)                             \
         PyUnicode_CheckExact(name) &&                           \
         PyUnicode_IS_READY(name) &&                             \
@@ -907,7 +908,7 @@ type_modified_unlocked(PyTypeObject *type)
     }
 
     type->tp_flags &= ~Py_TPFLAGS_VALID_VERSION_TAG;
-    type->tp_version_tag = 0; /* 0 is not a valid version tag */
+    FT_ATOMIC_STORE_UINT32_RELAXED(type->tp_version_tag, 0); /* 0 is not a valid version tag */
     if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
         // This field *must* be invalidated if the type is modified (see the
         // comment on struct _specialization_cache):
@@ -984,7 +985,7 @@ type_mro_modified(PyTypeObject *type, PyObject *bases) {
  clear:
     assert(!(type->tp_flags & _Py_TPFLAGS_STATIC_BUILTIN));
     type->tp_flags &= ~Py_TPFLAGS_VALID_VERSION_TAG;
-    type->tp_version_tag = 0; /* 0 is not a valid version tag */
+    FT_ATOMIC_STORE_UINT32_RELAXED(type->tp_version_tag, 0); /* 0 is not a valid version tag */
     if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
         // This field *must* be invalidated if the type is modified (see the
         // comment on struct _specialization_cache):
@@ -1020,7 +1021,8 @@ assign_version_tag(PyInterpreterState *interp, PyTypeObject *type)
             /* We have run out of version numbers */
             return 0;
         }
-        type->tp_version_tag = NEXT_GLOBAL_VERSION_TAG++;
+        FT_ATOMIC_STORE_UINT32_RELAXED(type->tp_version_tag,
+                                       NEXT_GLOBAL_VERSION_TAG++);
         assert (type->tp_version_tag <= _Py_MAX_GLOBAL_TYPE_VERSION_TAG);
     }
     else {
@@ -1029,7 +1031,8 @@ assign_version_tag(PyInterpreterState *interp, PyTypeObject *type)
             /* We have run out of version numbers */
             return 0;
         }
-        type->tp_version_tag = NEXT_VERSION_TAG(interp)++;
+        FT_ATOMIC_STORE_UINT32_RELAXED(type->tp_version_tag,
+                                       NEXT_VERSION_TAG(interp)++);
         assert (type->tp_version_tag != 0);
     }
 
@@ -5085,7 +5088,9 @@ _PyType_Lookup(PyTypeObject *type, PyObject *name)
     // synchronize-with other writing threads by doing an acquire load on the sequence
     while (1) {
         int sequence = _PySeqLock_BeginRead(&entry->sequence);
-        if (_Py_atomic_load_uint32_relaxed(&entry->version) == type->tp_version_tag &&
+        uint32_t entry_version = _Py_atomic_load_uint32_relaxed(&entry->version);
+        uint32_t type_version = _Py_atomic_load_uint32_relaxed(&type->tp_version_tag);
+        if (entry_version == type_version &&
             _Py_atomic_load_ptr_relaxed(&entry->name) == name) {
             assert(_PyType_HasFeature(type, Py_TPFLAGS_VALID_VERSION_TAG));
             OBJECT_STAT_INC_COND(type_cache_hits, !is_dunder_name(name));
