Make PyDictOrValues thread safe

Author: DinoV

Include/internal/pycore_dict.h

@@ -248,6 +248,7 @@ extern PyObject *_PyDict_FromItems(
         PyObject *const *keys, Py_ssize_t keys_offset,
         PyObject *const *values, Py_ssize_t values_offset,
         Py_ssize_t length);
+extern void _PyDict_FreeDictForDematerialization(PyDictObject *obj);
 
 static inline void
 _PyDictValues_AddToInsertionOrder(PyDictValues *values, Py_ssize_t ix)

Include/internal/pycore_lock.h

@@ -251,6 +251,7 @@ PyAPI_FUNC(void) _PyRWMutex_RUnlock(_PyRWMutex *rwmutex);
 PyAPI_FUNC(void) _PyRWMutex_Lock(_PyRWMutex *rwmutex);
 PyAPI_FUNC(void) _PyRWMutex_Unlock(_PyRWMutex *rwmutex);
 
+extern void _Py_yield(void);
 
 #ifdef __cplusplus
 }

Include/internal/pycore_object.h

@@ -631,6 +631,9 @@ typedef union {
     char *values;
 } PyDictOrValues;
 
+// Sentinel value to indicate when an update to PyDictOrValues is in-flight
+#define _PYDICTORVALUES_UPDATING 0x0000002
+
 static inline PyDictOrValues *
 _PyObject_DictOrValuesPointer(PyObject *obj)
 {
@@ -651,23 +654,120 @@ _PyDictOrValues_GetValues(PyDictOrValues dorv)
     return (PyDictValues *)(dorv.values + 1);
 }
 
+static inline PyDictValues *
+_PyDictOrValues_TryGetValues(PyDictOrValues *dorv)
+{
+#ifdef Py_GIL_DISABLED
+    char *values;
+    while (1) {
+        values = _Py_atomic_load_ptr_acquire(&dorv->values);
+        if (values != (char *)_PYDICTORVALUES_UPDATING) {
+            if ((uintptr_t)values & 1) {
+                return (PyDictValues *)(values + 1);
+            }
+            // The values have become a dict or is not yet initialized
+            return NULL;
+        }
+        // There is an atomic update of the values in progress...
+        _Py_yield();
+    }
+#else
+    if (_PyDictOrValues_IsValues(*dorv)) {
+        return _PyDictOrValues_GetValues(*dorv);
+    }
+    return NULL;
+#endif
+}
+
 static inline PyObject *
 _PyDictOrValues_GetDict(PyDictOrValues dorv)
 {
     assert(!_PyDictOrValues_IsValues(dorv));
     return dorv.dict;
 }
 
+// Trys to get the dict from the PyDictOrValues and returns
+// a new strong reference if successful.
+static inline PyObject *
+_PyDictOrValues_TryGetDict(PyDictOrValues *dorv)
+{
+#ifdef Py_GIL_DISABLED
+    PyObject *dict;
+    while (1) {
+        dict = _Py_atomic_load_ptr_acquire(&dorv->dict);
+        if (dict != (PyObject *)_PYDICTORVALUES_UPDATING) {
+            if ((uintptr_t)dict & 1) {
+                // The dict has become a values
+                return NULL;
+            } else if (dict != NULL) {
+                Py_INCREF(dict);
+                if (_Py_atomic_load_ptr_acquire(&dorv->dict) == dict) {
+                    return dict;
+                }
+
+                // We've lost a race (presumably with dematerialization) so
+                // we'll try again...
+                Py_DECREF(dict);
+            } else {
+                // The dict is not yet initialized...
+                return dict;
+            }
+        }
+        // There is an atomic update of the values in progress...
+        _Py_yield();
+    }
+#else
+    if (!_PyDictOrValues_IsValues(*dorv)) {
+        Py_XINCREF(dorv->dict);
+        return dorv->dict;
+    }
+    return NULL;
+#endif
+}
+
+static inline bool
+_PyDictOrValues_TrySetDict(PyDictOrValues *dorv_ptr, void *expected, PyObject *dict)
+{
+#ifdef Py_GIL_DISABLED
+    if (!_Py_atomic_compare_exchange_ptr(&dorv_ptr->dict, &expected, dict)) {
+        return false;
+    }
+    return true;
+#else
+    dorv_ptr->dict = dict;
+    return true;
+#endif
+}
+
+
+
+static inline void
+_PyDictOrValues_FreeValues(PyDictValues *values)
+{
+    // TODO: Maybe free with qsbr
+    int prefix_size = ((uint8_t *)values)[-1];
+    PyMem_Free(((char *)values)-prefix_size);
+}
+
 static inline void
 _PyDictOrValues_SetValues(PyDictOrValues *ptr, PyDictValues *values)
 {
+#ifdef Py_GIL_DISABLED
+    char *expected = NULL;
+    if (!_Py_atomic_compare_exchange_ptr(&ptr->values, &expected, ((char *)values) - 1)) {
+        _PyDictOrValues_FreeValues(values);
+    }
+#else
     ptr->values = ((char *)values) - 1;
+#endif
 }
 
 extern PyObject ** _PyObject_ComputedDictPointer(PyObject *);
 extern void _PyObject_FreeInstanceAttributes(PyObject *obj);
 extern int _PyObject_IsInstanceDictEmpty(PyObject *);
 
+extern int _PyObject_SetDict(PyObject *obj, PyObject *new_dict);
+
 // Export for 'math' shared extension
 PyAPI_FUNC(PyObject*) _PyObject_LookupSpecial(PyObject *, PyObject *);
 

Include/internal/pycore_opcode_metadata.h

@@ -114,16 +114,16 @@ const uint16_t _PyUop_Flags[MAX_UOP_ID+1] = {
     [_LOAD_ATTR] = HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG,
     [_GUARD_TYPE_VERSION] = HAS_DEOPT_FLAG | HAS_PASSTHROUGH_FLAG,
     [_CHECK_MANAGED_OBJECT_HAS_VALUES] = HAS_DEOPT_FLAG | HAS_PASSTHROUGH_FLAG,
-    [_LOAD_ATTR_INSTANCE_VALUE] = HAS_ARG_FLAG | HAS_DEOPT_FLAG,
+    [_LOAD_ATTR_INSTANCE_VALUE] = HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG,
     [_CHECK_ATTR_MODULE] = HAS_DEOPT_FLAG | HAS_PASSTHROUGH_FLAG,
     [_LOAD_ATTR_MODULE] = HAS_ARG_FLAG | HAS_DEOPT_FLAG,
-    [_CHECK_ATTR_WITH_HINT] = HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG | HAS_PASSTHROUGH_FLAG,
-    [_LOAD_ATTR_WITH_HINT] = HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG,
+    [_CHECK_ATTR_WITH_HINT] = HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG,
+    [_LOAD_ATTR_WITH_HINT] = HAS_ARG_FLAG | HAS_NAME_FLAG | HAS_DEOPT_FLAG,
     [_LOAD_ATTR_SLOT] = HAS_ARG_FLAG | HAS_DEOPT_FLAG,
     [_CHECK_ATTR_CLASS] = HAS_DEOPT_FLAG | HAS_PASSTHROUGH_FLAG,
     [_LOAD_ATTR_CLASS] = HAS_ARG_FLAG,
     [_GUARD_DORV_VALUES] = HAS_DEOPT_FLAG | HAS_PASSTHROUGH_FLAG,
-    [_STORE_ATTR_INSTANCE_VALUE] = HAS_ESCAPES_FLAG,
+    [_STORE_ATTR_INSTANCE_VALUE] = HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG,
     [_STORE_ATTR_SLOT] = HAS_ESCAPES_FLAG,
     [_COMPARE_OP] = HAS_ARG_FLAG | HAS_ERROR_FLAG | HAS_ESCAPES_FLAG,
     [_COMPARE_OP_FLOAT] = HAS_ARG_FLAG | HAS_DEOPT_FLAG | HAS_ESCAPES_FLAG,