Incref dict when it's usage is borrowed

DinoV · DinoV · commit 5a1b6fdfdaec · 2024-07-31T10:50:01.000-07:00
diff --git a/Lib/test/test_free_threading/test_dict.py b/Lib/test/test_free_threading/test_dict.py
@@ -142,6 +142,49 @@ def writer_func(l):
             for ref in thread_list:
                 self.assertIsNone(ref())
 
+    def test_racing_set_object_dict(self):
+        """Races assigning to __dict__ should be thread safe"""
+
+        class C: pass
+        f = C()
+        f.__dict__ = {"foo": 42}
+        THREAD_COUNT = 10
+        class MyDict(dict): pass
+
+        def writer_func():
+            for i in range(1000):
+                f.__dict__ = {"foo": 100}
+
+        def reader_func():
+            for i in range(1000):
+                f.foo
+
+        lists = []
+        readers = []
+        writers = []
+        for x in range(THREAD_COUNT):
+            thread_list = []
+            lists.append(thread_list)
+            writer = Thread(target=partial(writer_func))
+            writers.append(writer)
+
+        for x in range(THREAD_COUNT):
+            thread_list = []
+            lists.append(thread_list)
+            reader = Thread(target=partial(reader_func))
+            readers.append(reader)
+
+        for writer in writers:
+            writer.start()
+        for reader in readers:
+            reader.start()
+
+        for writer in writers:
+            writer.join()
+
+        for reader in readers:
+            reader.join()
+
     @unittest.skipIf(_testcapi is None, 'need _testcapi module')
     def test_dict_version(self):
         dict_version = _testcapi.dict_version
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -165,6 +165,8 @@ ASSERT_DICT_LOCKED(PyObject *op)
 
 #define IS_DICT_SHARED(mp) _PyObject_GC_IS_SHARED(mp)
 #define SET_DICT_SHARED(mp) _PyObject_GC_SET_SHARED(mp)
+#define IS_DICT_SHARED_INLINE(mp) _PyObject_GC_IS_SHARED_INLINE(mp)
+#define SET_DICT_SHARED_INLINE(mp) _PyObject_GC_SET_SHARED_INLINE(mp)
 #define LOAD_INDEX(keys, size, idx) _Py_atomic_load_int##size##_relaxed(&((const int##size##_t*)keys->dk_indices)[idx]);
 #define STORE_INDEX(keys, size, idx, value) _Py_atomic_store_int##size##_relaxed(&((int##size##_t*)keys->dk_indices)[idx], (int##size##_t)value);
 #define ASSERT_OWNED_OR_SHARED(mp) \
@@ -245,6 +247,8 @@ static inline void split_keys_entry_added(PyDictKeysObject *keys)
 #define UNLOCK_KEYS_IF_SPLIT(keys, kind)
 #define IS_DICT_SHARED(mp) (false)
 #define SET_DICT_SHARED(mp)
+#define IS_DICT_SHARED_INLINE(mp) (false)
+#define SET_DICT_SHARED_INLINE(mp)
 #define LOAD_INDEX(keys, size, idx) ((const int##size##_t*)(keys->dk_indices))[idx]
 #define STORE_INDEX(keys, size, idx, value) ((int##size##_t*)(keys->dk_indices))[idx] = (int##size##_t)value
 
@@ -3082,20 +3086,21 @@ dict_dealloc(PyObject *self)
     /* bpo-31095: UnTrack is needed before calling any callbacks */
     PyObject_GC_UnTrack(mp);
     Py_TRASHCAN_BEGIN(mp, dict_dealloc)
+    bool is_shared = IS_DICT_SHARED_INLINE(self);
     if (values != NULL) {
         if (values->embedded == 0) {
             for (i = 0, n = mp->ma_keys->dk_nentries; i < n; i++) {
                 Py_XDECREF(values->values[i]);
             }
-            free_values(values, false);
+            free_values(values, is_shared);
         }
-        dictkeys_decref(interp, keys, false);
+        dictkeys_decref(interp, keys, is_shared);
     }
     else if (keys != NULL) {
         assert(keys->dk_refcnt == 1 || keys == Py_EMPTY_KEYS);
-        dictkeys_decref(interp, keys, false);
+        dictkeys_decref(interp, keys, is_shared);
     }
-    if (Py_IS_TYPE(mp, &PyDict_Type)) {
+    if (Py_IS_TYPE(mp, &PyDict_Type) && !is_shared) {
         _Py_FREELIST_FREE(dicts, mp, Py_TYPE(mp)->tp_free);
     }
     else {
@@ -7047,6 +7052,9 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
             FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
                                 (PyDictObject *)Py_XNewRef(new_dict));
         }
+        // We could be racing with a read from a borrowed reference of
+        // the dict, so we need to free it using QSBR.
+        SET_DICT_SHARED_INLINE(dict);
         Py_END_CRITICAL_SECTION2();
 
         if (err == 0) {
