From 5da0280648b5f1c5142dad39dcd1e7fd99fdf37a Mon Sep 17 00:00:00 2001
From: Dino Viehland <dinoviehland@meta.com>
Date: Thu, 25 Apr 2024 08:53:29 -0700
Subject: [PATCH] gh-117657: Fixes a few small TSAN issues in dictobject
 (#118200)

Fixup TSAN errors for dict
---
 Include/cpython/dictobject.h               |  4 ++++
 Include/internal/pycore_object.h           |  2 +-
 Objects/dictobject.c                       | 25 ++++++++++++----------
 Tools/tsan/suppressions_free_threading.txt |  3 ---
 4 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/Include/cpython/dictobject.h b/Include/cpython/dictobject.h
index 35b6a822a0d..3fd23b9313c 100644
--- a/Include/cpython/dictobject.h
+++ b/Include/cpython/dictobject.h
@@ -56,7 +56,11 @@ static inline Py_ssize_t PyDict_GET_SIZE(PyObject *op) {
     PyDictObject *mp;
     assert(PyDict_Check(op));
     mp = _Py_CAST(PyDictObject*, op);
+#ifdef Py_GIL_DISABLED
+    return _Py_atomic_load_ssize_relaxed(&mp->ma_used);
+#else
     return mp->ma_used;
+#endif
 }
 #define PyDict_GET_SIZE(op) PyDict_GET_SIZE(_PyObject_CAST(op))
 
diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
index 88b052f4544..7df8003196d 100644
--- a/Include/internal/pycore_object.h
+++ b/Include/internal/pycore_object.h
@@ -688,7 +688,7 @@ static inline PyDictObject *
 _PyObject_GetManagedDict(PyObject *obj)
 {
     PyManagedDictPointer *dorv = _PyObject_ManagedDictPointer(obj);
-    return (PyDictObject *)FT_ATOMIC_LOAD_PTR_RELAXED(dorv->dict);
+    return (PyDictObject *)FT_ATOMIC_LOAD_PTR_ACQUIRE(dorv->dict);
 }
 
 static inline PyDictValues *
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index 2644516bc30..afcf535f8c0 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -1097,10 +1097,11 @@ compare_unicode_unicode(PyDictObject *mp, PyDictKeysObject *dk,
                         void *ep0, Py_ssize_t ix, PyObject *key, Py_hash_t hash)
 {
     PyDictUnicodeEntry *ep = &((PyDictUnicodeEntry *)ep0)[ix];
-    assert(ep->me_key != NULL);
-    assert(PyUnicode_CheckExact(ep->me_key));
-    if (ep->me_key == key ||
-            (unicode_get_hash(ep->me_key) == hash && unicode_eq(ep->me_key, key))) {
+    PyObject *ep_key = FT_ATOMIC_LOAD_PTR_RELAXED(ep->me_key);
+    assert(ep_key != NULL);
+    assert(PyUnicode_CheckExact(ep_key));
+    if (ep_key == key ||
+            (unicode_get_hash(ep_key) == hash && unicode_eq(ep_key, key))) {
         return 1;
     }
     return 0;
@@ -1761,10 +1762,12 @@ insertdict(PyInterpreterState *interp, PyDictObject *mp,
         else {
             assert(old_value != NULL);
             if (DK_IS_UNICODE(mp->ma_keys)) {
-                DK_UNICODE_ENTRIES(mp->ma_keys)[ix].me_value = value;
+                PyDictUnicodeEntry *ep = &DK_UNICODE_ENTRIES(mp->ma_keys)[ix];
+                STORE_VALUE(ep, value);
             }
             else {
-                DK_ENTRIES(mp->ma_keys)[ix].me_value = value;
+                PyDictKeyEntry *ep = &DK_ENTRIES(mp->ma_keys)[ix];
+                STORE_VALUE(ep, value);
             }
         }
         mp->ma_version_tag = new_version;
@@ -1810,15 +1813,15 @@ insert_to_emptydict(PyInterpreterState *interp, PyDictObject *mp,
     if (unicode) {
         PyDictUnicodeEntry *ep = DK_UNICODE_ENTRIES(newkeys);
         ep->me_key = key;
-        ep->me_value = value;
+        STORE_VALUE(ep, value);
     }
     else {
         PyDictKeyEntry *ep = DK_ENTRIES(newkeys);
         ep->me_key = key;
         ep->me_hash = hash;
-        ep->me_value = value;
+        STORE_VALUE(ep, value);
     }
-    mp->ma_used++;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(mp->ma_used, FT_ATOMIC_LOAD_SSIZE_RELAXED(mp->ma_used) + 1);
     mp->ma_version_tag = new_version;
     newkeys->dk_usable--;
     newkeys->dk_nentries++;
@@ -2510,7 +2513,7 @@ delitem_common(PyDictObject *mp, Py_hash_t hash, Py_ssize_t ix,
     Py_ssize_t hashpos = lookdict_index(mp->ma_keys, hash, ix);
     assert(hashpos >= 0);
 
-    mp->ma_used--;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(mp->ma_used, FT_ATOMIC_LOAD_SSIZE(mp->ma_used) - 1);
     mp->ma_version_tag = new_version;
     if (_PyDict_HasSplitTable(mp)) {
         assert(old_value == mp->ma_values->values[ix]);
@@ -6895,7 +6898,7 @@ _PyObject_TryGetInstanceAttribute(PyObject *obj, PyObject *name, PyObject **attr
     }
 
 #ifdef Py_GIL_DISABLED
-    PyObject *value = _Py_atomic_load_ptr_relaxed(&values->values[ix]);
+    PyObject *value = _Py_atomic_load_ptr_acquire(&values->values[ix]);
     if (value == NULL || _Py_TryIncrefCompare(&values->values[ix], value)) {
         *attr = value;
         return true;
diff --git a/Tools/tsan/suppressions_free_threading.txt b/Tools/tsan/suppressions_free_threading.txt
index aa954ca9b6c..e4ca32bebc5 100644
--- a/Tools/tsan/suppressions_free_threading.txt
+++ b/Tools/tsan/suppressions_free_threading.txt
@@ -28,9 +28,6 @@ race:_PyObject_GC_TRACK
 race:_PyParkingLot_Park
 race:_PyType_HasFeature
 race:assign_version_tag
-race:compare_unicode_unicode
-race:delitem_common
-race:dictresize
 race:gc_collect_main
 race:gc_restore_tid
 race:initialize_new_array