Incref dict when it's usage is borrowed

Author: DinoV

Include/internal/pycore_gc.h

@@ -130,11 +130,7 @@ static inline int _PyObject_GC_IS_SHARED_INLINE(PyObject *op) {
 #define _PyObject_GC_IS_SHARED_INLINE(op) \
     _PyObject_GC_IS_SHARED_INLINE(_Py_CAST(PyObject*, op))
 
-static inline void _PyObject_GC_SET_SHARED_INLINE(PyObject *op) {
-    _PyObject_SET_GC_BITS(op, _PyGC_BITS_SHARED_INLINE);
-}
-#define _PyObject_GC_SET_SHARED_INLINE(op) \
-    _PyObject_GC_SET_SHARED_INLINE(_Py_CAST(PyObject*, op))
+extern void _PyObject_GC_SET_SHARED_INLINE(PyObject *op);
 
 #endif
 

Include/refcount.h

@@ -277,6 +277,10 @@ PyAPI_FUNC(void) _Py_DecRefSharedDebug(PyObject *, const char *, int);
 // zero. Otherwise, the thread gives up ownership and merges the reference
 // count fields.
 PyAPI_FUNC(void) _Py_MergeZeroLocalRefcount(PyObject *);
+
+// Queue an object to be deferred at the next quiescent state
+PyAPI_FUNC(void) _PyObject_FreeDeferred(PyObject *o);
+
 #endif
 
 #if defined(Py_LIMITED_API) && (Py_LIMITED_API+0 >= 0x030c0000 || defined(Py_REF_DEBUG))

Lib/test/test_free_threading/test_dict.py

@@ -142,6 +142,70 @@ 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
+        class MyDict(dict): pass
+        for cyclic in (False, True):
+            f = C()
+            f.__dict__ = {"foo": 42}
+            THREAD_COUNT = 10
+
+            def writer_func(l):
+                for i in range(100000):
+                    if cyclic:
+                        other_d = {}
+                    d = MyDict({"foo": 100})
+                    if cyclic:
+                        d["x"] = other_d
+                        other_d["bar"] = d
+                    l.append(weakref.ref(d))
+                    f.__dict__ = d
+
+            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, thread_list))
+                writers.append(writer)
+
+            for x in range(THREAD_COUNT):
+                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()
+
+            f.__dict__ = {}
+            gc.collect()
+            gc.collect()
+
+            count = 0
+            ids = set()
+            for thread_list in lists:
+                for i, ref in enumerate(thread_list):
+                    if ref() is None:
+                        continue
+                    count += 1
+                    ids.add(id(ref()))
+                    count += 1
+
+            self.assertEqual(count, 0)
+
     @unittest.skipIf(_testcapi is None, 'need _testcapi module')
     def test_dict_version(self):
         dict_version = _testcapi.dict_version

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
 
@@ -3198,7 +3202,7 @@ dict_dealloc(PyObject *self)
         assert(keys->dk_refcnt == 1 || keys == Py_EMPTY_KEYS);
         dictkeys_decref(interp, keys, false);
     }
-    if (Py_IS_TYPE(mp, &PyDict_Type)) {
+    if (Py_IS_TYPE(mp, &PyDict_Type) && !IS_DICT_SHARED_INLINE(mp)) {
         _Py_FREELIST_FREE(dicts, mp, Py_TYPE(mp)->tp_free);
     }
     else {
@@ -7126,6 +7130,75 @@ set_dict_inline_values(PyObject *obj, PyDictObject *new_dict)
     }
 }
 
+#ifdef Py_GIL_DISABLED
+
+// Trys and sets the dictionary for an object in the easy case when our current
+// dictionary is either completely not materialized or is a dictionary which
+// does not point at the inline values.
+static bool
+try_set_dict_inline_only_or_other_dict(PyObject *obj, PyObject *new_dict, PyDictObject **cur_dict)
+{
+    bool replaced = false;
+    Py_BEGIN_CRITICAL_SECTION(obj);
+
+    PyDictObject *dict = *cur_dict = _PyObject_GetManagedDict(obj);
+    if (dict == NULL) {
+        // We only have inline values, we can just completely replace them.
+        set_dict_inline_values(obj, (PyDictObject *)new_dict);
+        replaced = true;
+        goto exit_lock;
+    }
+
+    if (FT_ATOMIC_LOAD_PTR_RELAXED(dict->ma_values) != _PyObject_InlineValues(obj)) {
+        // We have a materialized dict which doesn't point at the inline values,
+        // We get to simply swap dictionaries and free the old dictionary.
+        FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
+                (PyDictObject *)Py_XNewRef(new_dict));
+        replaced = true;
+        goto exit_lock;
+    } else {
+        // We have inline values, we need to lock the dict and the object
+        // at the same time to safely dematerialize them. To do that while releasing
+        // the object lock we need a strong reference to the current dictionary.
+        Py_INCREF(dict);
+    }
+exit_lock:
+    Py_END_CRITICAL_SECTION();
+    return replaced;
+}
+
+#endif
+
+// Replaces a dictionary that is probably the dictionary which has been
+// materialized and points at the inline values. We could have raced
+// and replaced it with another dictionary though.
+static int
+replace_dict_probably_inline_materialized(PyObject *obj, PyDictObject *inline_dict,
+                                          PyObject *new_dict, PyDictObject **replaced_dict)
+{
+    // But we could have had another thread race in after we released
+    // the object lock
+    int err = 0;
+    *replaced_dict = _PyObject_GetManagedDict(obj);
+    assert(FT_ATOMIC_LOAD_PTR_RELAXED(inline_dict->ma_values) == _PyObject_InlineValues(obj));
+
+    if (*replaced_dict == inline_dict) {
+        err = _PyDict_DetachFromObject(inline_dict, obj);
+        if (err != 0) {
+            return err;
+        }
+        SET_DICT_SHARED_INLINE((PyObject *)inline_dict);
+        // We incref'd the inline dict and the object owns a ref.
+        // Clear the object's reference, we'll clear the local
+        // reference after releasing the lock.
+        Py_CLEAR(*replaced_dict);
+    }
+
+    FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
+                    (PyDictObject *)Py_XNewRef(new_dict));
+    return err;
+}
+
 int
 _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
 {
@@ -7134,42 +7207,51 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
     int err = 0;
     PyTypeObject *tp = Py_TYPE(obj);
     if (tp->tp_flags & Py_TPFLAGS_INLINE_VALUES) {
-        PyDictObject *dict = _PyObject_GetManagedDict(obj);
-        if (dict == NULL) {
 #ifdef Py_GIL_DISABLED
-            Py_BEGIN_CRITICAL_SECTION(obj);
+        PyDictObject *prev_dict;
+        if (!try_set_dict_inline_only_or_other_dict(obj, new_dict, &prev_dict)) {
+            // We had a materialized dictionary which pointed at the inline
+            // values. We need to lock both the object and the dict at the
+            // same time to safely replace it. We can't merely lock the dictionary
+            // while the object is locked because it could suspend the object lock.
+            PyDictObject *replaced_dict;
 
-            dict = _PyObject_ManagedDictPointer(obj)->dict;
-            if (dict == NULL) {
-                set_dict_inline_values(obj, (PyDictObject *)new_dict);
-            }
+            assert(prev_dict != NULL);
+            Py_BEGIN_CRITICAL_SECTION2(obj, prev_dict);
 
-            Py_END_CRITICAL_SECTION();
+            err = replace_dict_probably_inline_materialized(obj, prev_dict, new_dict, &replaced_dict);
 
-            if (dict == NULL) {
-                return 0;
+            Py_END_CRITICAL_SECTION2();
+
+            Py_DECREF(prev_dict);
+            if (err != 0) {
+                return err;
             }
+            prev_dict = replaced_dict;
+        }
+
+        if (prev_dict != NULL) {
+            Py_BEGIN_CRITICAL_SECTION(prev_dict);
+            SET_DICT_SHARED_INLINE((PyObject *)prev_dict);
+            Py_END_CRITICAL_SECTION();
+            // Readers from the old dictionary use a borrowed reference. We need
+            // to set the dict to be freed via QSBR which requires locking it.
+            Py_DECREF(prev_dict);
+        }
+        return 0;
 #else
+        PyDictObject *dict = _PyObject_GetManagedDict(obj);
+        if (dict == NULL) {
             set_dict_inline_values(obj, (PyDictObject *)new_dict);
             return 0;
-#endif
-        }
-
-        Py_BEGIN_CRITICAL_SECTION2(dict, obj);
-
-        // We've locked dict, but the actual dict could have changed
-        // since we locked it.
-        dict = _PyObject_ManagedDictPointer(obj)->dict;
-        err = _PyDict_DetachFromObject(dict, obj);
-        if (err == 0) {
-            FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
-                                (PyDictObject *)Py_XNewRef(new_dict));
         }
-        Py_END_CRITICAL_SECTION2();
-
-        if (err == 0) {
-            Py_XDECREF(dict);
+        if (_PyDict_DetachFromObject(dict, obj) == 0) {
+            _PyObject_ManagedDictPointer(obj)->dict = (PyDictObject *)Py_XNewRef(new_dict);
+            Py_DECREF(dict);
+            return 0;
         }
+        return -1;
+#endif
     }
     else {
         PyDictObject *dict;
@@ -7182,7 +7264,13 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
                             (PyDictObject *)Py_XNewRef(new_dict));
 
         Py_END_CRITICAL_SECTION();
-
+#ifdef Py_GIL_DISABLED
+        if (dict != NULL) {
+            Py_BEGIN_CRITICAL_SECTION(dict);
+            SET_DICT_SHARED_INLINE((PyObject *)dict);
+            Py_END_CRITICAL_SECTION();
+        }
+#endif
         Py_XDECREF(dict);
     }
     assert(_PyObject_InlineValuesConsistencyCheck(obj));

Objects/object.c

@@ -319,6 +319,27 @@ is_dead(PyObject *o)
 }
 # endif
 
+void
+_PyObject_FreeDeferred(PyObject *o)
+{
+    // It's possible that before we've hit the quiescent-state that it
+    // has been incref'd and then decref'd back to zero again. We re-use
+    // the finalized bit to track if we've already queued the delay freed
+    // for this object as we no longer need to GC track it if it's being
+    // freed by QSBR as it's reclaimed by the ref count.
+    if (!_PyGC_FINALIZED(o)) {
+        Py_BEGIN_CRITICAL_SECTION(o);
+        if (!_PyGC_FINALIZED(o)) {
+            if (_PyObject_GC_IS_TRACKED(o)) {
+                PyObject_GC_UnTrack(o);
+            }
+            _PyGC_SET_FINALIZED(o);
+            _PyObject_FreeDelayed(o);
+        }
+        Py_END_CRITICAL_SECTION();
+    }
+}
+
 void
 _Py_DecRefSharedDebug(PyObject *o, const char *filename, int lineno)
 {
@@ -360,8 +381,12 @@ _Py_DecRefSharedDebug(PyObject *o, const char *filename, int lineno)
         _Py_brc_queue_object(o);
     }
     else if (new_shared == _Py_REF_MERGED) {
-        // refcount is zero AND merged
-        _Py_Dealloc(o);
+        if (!_PyObject_GC_IS_SHARED_INLINE(o)) {
+            // refcount is zero AND merged
+            _Py_Dealloc(o);
+        } else {
+            _PyObject_FreeDeferred(o);
+        }
     }
 }
 
@@ -379,7 +404,12 @@ _Py_MergeZeroLocalRefcount(PyObject *op)
     Py_ssize_t shared = _Py_atomic_load_ssize_acquire(&op->ob_ref_shared);
     if (shared == 0) {
         // Fast-path: shared refcount is zero (including flags)
-        _Py_Dealloc(op);
+        if (!_PyObject_GC_IS_SHARED_INLINE(op)) {
+            // refcount is zero AND merged
+            _Py_Dealloc(op);
+        } else {
+            _PyObject_FreeDeferred(op);
+        }
         return;
     }
 
@@ -398,7 +428,12 @@ _Py_MergeZeroLocalRefcount(PyObject *op)
     if (new_shared == _Py_REF_MERGED) {
         // i.e., the shared refcount is zero (only the flags are set) so we
         // deallocate the object.
-        _Py_Dealloc(op);
+        if (!_PyObject_GC_IS_SHARED_INLINE(op)) {
+            // refcount is zero AND merged
+            _Py_Dealloc(op);
+        } else {
+            _PyObject_FreeDeferred(op);
+        }
     }
 }
 

Objects/obmalloc.c

@@ -1096,7 +1096,7 @@ static void
 free_work_item(uintptr_t ptr)
 {
     if (ptr & 0x01) {
-        PyObject_Free((char *)(ptr - 1));
+        _Py_Dealloc((PyObject*)(char *)(ptr - 1));
     }
     else {
         PyMem_Free((void *)ptr);

Python/brc.c

@@ -107,7 +107,11 @@ merge_queued_objects(_PyObjectStack *to_merge)
         // Subtract one when merging because the queue had a reference.
         Py_ssize_t refcount = _Py_ExplicitMergeRefcount(ob, -1);
         if (refcount == 0) {
-            _Py_Dealloc(ob);
+            if (!_PyObject_GC_IS_SHARED_INLINE(ob)) {
+                _Py_Dealloc(ob);
+            } else {
+                _PyObject_FreeDeferred(ob);
+            }
         }
     }
 }