dict: make dict thread-safe

Include/cpython/dictobject.h

@@ -26,6 +26,10 @@ typedef struct {
        If ma_values is not NULL, the table is split:
        keys are stored in ma_keys and values are stored in ma_values */
     PyDictValues *ma_values;
+
+    PyMutex ma_mutex;
+
+    uint8_t ma_maybe_shared;
 } PyDictObject;
 
 PyAPI_FUNC(PyObject *) _PyDict_GetItem_KnownHash(PyObject *mp, PyObject *key,
@@ -36,12 +40,15 @@ PyAPI_FUNC(PyObject *) _PyDict_GetItemIdWithError(PyObject *dp,
 PyAPI_FUNC(PyObject *) _PyDict_GetItemStringWithError(PyObject *, const char *);
 PyAPI_FUNC(PyObject *) PyDict_SetDefault(
     PyObject *mp, PyObject *key, PyObject *defaultobj);
+PyAPI_FUNC(PyObject *) _PyDict_SetDefault(PyObject *d, PyObject *key, PyObject *defaultobj,
+                                          int incref, int *is_insert);
 PyAPI_FUNC(int) _PyDict_SetItem_KnownHash(PyObject *mp, PyObject *key,
                                           PyObject *item, Py_hash_t hash);
 PyAPI_FUNC(int) _PyDict_DelItem_KnownHash(PyObject *mp, PyObject *key,
                                           Py_hash_t hash);
 PyAPI_FUNC(int) _PyDict_DelItemIf(PyObject *mp, PyObject *key,
-                                  int (*predicate)(PyObject *value));
+                                  int (*predicate)(PyObject *value, void *data),
+                                  void *data);
 PyAPI_FUNC(int) _PyDict_Next(
     PyObject *mp, Py_ssize_t *pos, PyObject **key, PyObject **value, Py_hash_t *hash);
 

Include/cpython/pystate.h

@@ -106,6 +106,10 @@ typedef struct _stack_chunk {
     PyObject * data[1]; /* Variable sized */
 } _PyStackChunk;
 
+typedef struct _Py_dict_thread_state {
+    uint64_t dict_version;
+} _Py_dict_thread_state;
+
 struct mi_heap_s;
 typedef struct mi_heap_s mi_heap_t;
 typedef struct _PyEventRc _PyEventRc;
@@ -150,6 +154,8 @@ struct _ts {
      * or most recently, executing _PyEval_EvalFrameDefault. */
     _PyCFrame *cframe;
 
+    _Py_dict_thread_state dict_state;
+
     /* The thread will not stop for GC or other stop-the-world requests.
      * Used for *short* critical sections that to prevent deadlocks between
      * finalizers and stopped threads. */

Include/dictobject.h

@@ -19,6 +19,9 @@ PyAPI_DATA(PyTypeObject) PyDict_Type;
 #define PyDict_CheckExact(op) Py_IS_TYPE((op), &PyDict_Type)
 
 PyAPI_FUNC(PyObject *) PyDict_New(void);
+PyAPI_FUNC(PyObject *) PyDict_FetchItem(PyObject *mp, PyObject *key);
+PyAPI_FUNC(PyObject *) PyDict_FetchItemString(PyObject *dp, const char *key);
+PyAPI_FUNC(PyObject *) PyDict_FetchItemWithError(PyObject *mp, PyObject *key);
 PyAPI_FUNC(PyObject *) PyDict_GetItem(PyObject *mp, PyObject *key);
 PyAPI_FUNC(PyObject *) PyDict_GetItemWithError(PyObject *mp, PyObject *key);
 PyAPI_FUNC(int) PyDict_SetItem(PyObject *mp, PyObject *key, PyObject *item);
@@ -67,9 +70,9 @@ PyAPI_DATA(PyTypeObject) PyDictKeys_Type;
 PyAPI_DATA(PyTypeObject) PyDictValues_Type;
 PyAPI_DATA(PyTypeObject) PyDictItems_Type;
 
-#define PyDictKeys_Check(op) PyObject_TypeCheck((op), &PyDictKeys_Type)
-#define PyDictValues_Check(op) PyObject_TypeCheck((op), &PyDictValues_Type)
-#define PyDictItems_Check(op) PyObject_TypeCheck((op), &PyDictItems_Type)
+#define PyDictKeys_Check(op) Py_IS_TYPE((op), &PyDictKeys_Type)
+#define PyDictValues_Check(op) Py_IS_TYPE((op), &PyDictValues_Type)
+#define PyDictItems_Check(op) Py_IS_TYPE((op), &PyDictItems_Type)
 /* This excludes Values, since they are not sets. */
 # define PyDictViewSet_Check(op) \
     (PyDictKeys_Check(op) || PyDictItems_Check(op))

Include/internal/pycore_dict.h

@@ -11,6 +11,7 @@ extern "C" {
 
 #include "pycore_dict_state.h"
 #include "pycore_runtime.h"         // _PyRuntime
+#include "pycore_pystate.h"         // _PyThreadState_GET()
 
 
 /* runtime lifecycle */
@@ -22,9 +23,9 @@ extern void _PyDict_Fini(PyInterpreterState *interp);
 
 typedef struct {
     /* Cached hash code of me_key. */
-    Py_hash_t me_hash;
     PyObject *me_key;
     PyObject *me_value; /* This field is only meaningful for combined tables */
+    Py_hash_t me_hash;
 } PyDictKeyEntry;
 
 typedef struct {
@@ -64,12 +65,13 @@ extern PyObject *_PyDict_Pop_KnownHash(PyObject *, PyObject *, Py_hash_t, PyObje
 typedef enum {
     DICT_KEYS_GENERAL = 0,
     DICT_KEYS_UNICODE = 1,
-    DICT_KEYS_SPLIT = 2
+    DICT_KEYS_SPLIT = 2,
 } DictKeysKind;
 
 /* See dictobject.c for actual layout of DictKeysObject */
 struct _dictkeysobject {
-    Py_ssize_t dk_refcnt;
+    /* Mutex to prevent concurrent modification of shared keys. */
+    _PyMutex dk_mutex;
 
     /* Size of the hash table (dk_indices). It must be a power of 2. */
     uint8_t dk_log2_size;
@@ -108,6 +110,13 @@ struct _dictkeysobject {
        see the DK_ENTRIES() macro */
 };
 
+typedef struct PyDictSharedKeysObject {
+    uint8_t tracked;
+    uint8_t marked;
+    struct PyDictSharedKeysObject *next;
+    struct _dictkeysobject keys;
+} PyDictSharedKeysObject;
+
 /* This must be no more than 250, for the prefix size to fit in one byte. */
 #define SHARED_KEYS_MAX_SIZE 30
 #define NEXT_LOG2_SHARED_KEYS_MAX_SIZE 6
@@ -142,14 +151,31 @@ static inline PyDictUnicodeEntry* DK_UNICODE_ENTRIES(PyDictKeysObject *dk) {
     assert(dk->dk_kind != DICT_KEYS_GENERAL);
     return (PyDictUnicodeEntry*)_DK_ENTRIES(dk);
 }
+static inline PyDictSharedKeysObject* DK_AS_SPLIT(PyDictKeysObject *dk) {
+    assert(dk->dk_kind == DICT_KEYS_SPLIT);
+    char *mem = (char *)dk - offsetof(PyDictSharedKeysObject, keys);
+    return (PyDictSharedKeysObject *)mem;
+}
 
 #define DK_IS_UNICODE(dk) ((dk)->dk_kind != DICT_KEYS_GENERAL)
 
 #define DICT_VERSION_INCREMENT (1 << DICT_MAX_WATCHERS)
+#define DICT_GLOBAL_VERSION_INCREMENT (DICT_VERSION_INCREMENT * 256)
 #define DICT_VERSION_MASK (DICT_VERSION_INCREMENT - 1)
 
-#define DICT_NEXT_VERSION() \
-    (_PyRuntime.dict_state.global_version += DICT_VERSION_INCREMENT)
+static inline uint64_t
+_PyDict_NextVersion(PyThreadState *tstate)
+{
+    uint64_t version = tstate->dict_state.dict_version;
+    if (version % DICT_GLOBAL_VERSION_INCREMENT == 0) {
+        version = _Py_atomic_add_uint64(
+            &_PyRuntime.dict_state.global_version,
+            DICT_GLOBAL_VERSION_INCREMENT);
+    }
+    version += DICT_VERSION_INCREMENT;
+    tstate->dict_state.dict_version = version;
+    return version;
+}
 
 void
 _PyDict_SendEvent(int watcher_bits,
@@ -167,9 +193,9 @@ _PyDict_NotifyEvent(PyDict_WatchEvent event,
     int watcher_bits = mp->ma_version_tag & DICT_VERSION_MASK;
     if (watcher_bits) {
         _PyDict_SendEvent(watcher_bits, event, mp, key, value);
-        return DICT_NEXT_VERSION() | watcher_bits;
+        return _PyDict_NextVersion(_PyThreadState_GET()) | watcher_bits;
     }
-    return DICT_NEXT_VERSION();
+    return _PyDict_NextVersion(_PyThreadState_GET());
 }
 
 extern PyObject *_PyObject_MakeDictFromInstanceAttributes(PyObject *obj, PyDictValues *values);
@@ -184,7 +210,7 @@ _PyDictValues_AddToInsertionOrder(PyDictValues *values, Py_ssize_t ix)
     assert(ix < SHARED_KEYS_MAX_SIZE);
     uint8_t *size_ptr = ((uint8_t *)values)-2;
     int size = *size_ptr;
-    assert(size+2 < ((uint8_t *)values)[-1]);
+    assert(size < ((uint8_t *)values)[-1]);
     size++;
     size_ptr[-size] = (uint8_t)ix;
     *size_ptr = size;

Include/internal/pycore_dict_state.h

@@ -29,6 +29,8 @@ struct _Py_dict_runtime_state {
 
 #define DICT_MAX_WATCHERS 8
 
+typedef struct PyDictSharedKeysObject PyDictSharedKeysObject;
+
 struct _Py_dict_state {
 #if PyDict_MAXFREELIST > 0
     /* Dictionary reuse scheme to save calls to malloc and free */
@@ -38,6 +40,8 @@ struct _Py_dict_state {
     int keys_numfree;
 #endif
     PyDict_WatchCallback watchers[DICT_MAX_WATCHERS];
+    /* shared keys from deallocated types (i.e., potentially dead) */
+    PyDictSharedKeysObject *tracked_shared_keys;
 };
 
 

Include/internal/pycore_object.h

@@ -489,17 +489,25 @@ _PyObject_DictOrValuesPointer(PyObject *obj)
     return (PyDictOrValues *)((char *)obj + MANAGED_DICT_OFFSET);
 }
 
+static inline PyDictOrValues
+_PyObject_DictOrValues(PyObject *obj)
+{
+    PyDictOrValues dorv;
+    dorv.values = _Py_atomic_load_ptr_relaxed(_PyObject_DictOrValuesPointer(obj));
+    return dorv;
+}
+
 static inline int
 _PyDictOrValues_IsValues(PyDictOrValues dorv)
 {
-    return ((uintptr_t)dorv.values) & 1;
+    return ((uintptr_t)dorv.values & 4) != 0;
 }
 
 static inline PyDictValues *
 _PyDictOrValues_GetValues(PyDictOrValues dorv)
 {
     assert(_PyDictOrValues_IsValues(dorv));
-    return (PyDictValues *)(dorv.values + 1);
+    return (PyDictValues *)((uintptr_t)dorv.values & ~7);
 }
 
 static inline PyObject *
@@ -512,7 +520,47 @@ _PyDictOrValues_GetDict(PyDictOrValues dorv)
 static inline void
 _PyDictOrValues_SetValues(PyDictOrValues *ptr, PyDictValues *values)
 {
-    ptr->values = ((char *)values) - 1;
+    ptr->values = ((char *)values) + 4;
+}
+
+extern PyDictValues*
+_PyDictValues_LockSlow(PyDictOrValues *dorv_ptr);
+
+extern void
+_PyDictValues_UnlockSlow(PyDictOrValues *dorv_ptr);
+
+extern void
+_PyDictValues_UnlockDict(PyDictOrValues *dorv_ptr, PyObject *dict);
+
+static inline PyDictValues *
+_PyDictValues_Lock(PyDictOrValues *dorv_ptr)
+{
+    PyDictOrValues dorv;
+    dorv.values = _Py_atomic_load_ptr_relaxed(dorv_ptr);
+    if (!_PyDictOrValues_IsValues(dorv)) {
+        return NULL;
+    }
+    uintptr_t v = (uintptr_t)dorv.values;
+    if ((v & LOCKED) == UNLOCKED) {
+        if (_Py_atomic_compare_exchange_ptr(dorv_ptr, dorv.values, dorv.values + LOCKED)) {
+            return _PyDictOrValues_GetValues(dorv);
+        }
+    }
+    return _PyDictValues_LockSlow(dorv_ptr);
+}
+
+static inline void
+_PyDictValues_Unlock(PyDictOrValues *dorv_ptr)
+{
+    char *values = _Py_atomic_load_ptr_relaxed(&dorv_ptr->values);
+    uintptr_t v = (uintptr_t)values;
+    assert((v & LOCKED));
+    if ((v & HAS_PARKED) == 0) {
+        if (_Py_atomic_compare_exchange_ptr(dorv_ptr, values, values - LOCKED)) {
+            return;
+        }
+    }
+    _PyDictValues_UnlockSlow(dorv_ptr);
 }
 
 extern PyObject ** _PyObject_ComputedDictPointer(PyObject *);

Lib/test/support/__init__.py

@@ -1822,19 +1822,23 @@ def _check_tracemalloc():
 
 
 def check_free_after_iterating(test, iter, cls, args=()):
-    class A(cls):
-        def __del__(self):
-            nonlocal done
-            done = True
-            try:
-                next(it)
-            except StopIteration:
-                pass
-
     done = False
-    it = iter(A(*args))
-    # Issue 26494: Shouldn't crash
-    test.assertRaises(StopIteration, next, it)
+
+    def wrapper():
+        class A(cls):
+            def __del__(self):
+                nonlocal done
+                done = True
+                try:
+                    next(it)
+                except StopIteration:
+                    pass
+
+        it = iter(A(*args))
+        # Issue 26494: Shouldn't crash
+        test.assertRaises(StopIteration, next, it)
+
+    wrapper()
     # The sequence should be deallocated just after the end of iterating
     gc_collect()
     test.assertTrue(done)

Misc/stable_abi.toml

@@ -2225,6 +2225,15 @@
     added = '3.12'
     abi_only = true
 
+# New reference versions of existing APIs
+
+[function.PyDict_FetchItem]
+    added = '3.12'
+[function.PyDict_FetchItemString]
+    added = '3.12'
+[function.PyDict_FetchItemWithError]
+    added = '3.12'
+
 # Support for Stable ABI in debug builds
 
 [function._Py_NegativeRefcount]

Modules/_weakref.c

@@ -38,7 +38,7 @@ _weakref_getweakrefcount_impl(PyObject *module, PyObject *object)
 
 
 static int
-is_dead_weakref(PyObject *value)
+is_dead_weakref(PyObject *value, void *_unused)
 {
     int is_dead;
     if (!PyWeakref_Check(value)) {
@@ -68,7 +68,7 @@ _weakref__remove_dead_weakref_impl(PyObject *module, PyObject *dct,
                                    PyObject *key)
 /*[clinic end generated code: output=d9ff53061fcb875c input=19fc91f257f96a1d]*/
 {
-    if (_PyDict_DelItemIf(dct, key, is_dead_weakref) < 0) {
+    if (_PyDict_DelItemIf(dct, key, is_dead_weakref, NULL) < 0) {
         if (PyErr_ExceptionMatches(PyExc_KeyError))
             /* This function is meant to allow safe weak-value dicts
                with GC in another thread (see issue #28427), so it's