gh-112069: Make _PySet_NextEntry to be thread-safe.

Author: corona10

Include/internal/pycore_setobject.h

@@ -15,6 +15,13 @@ PyAPI_FUNC(int) _PySet_NextEntry(
     PyObject **key,
     Py_hash_t *hash);
 
+// Export for 'Python/compile.c'
+PyAPI_FUNC(int) _PyFrozenSet_NextEntry(
+    PyObject *set,
+    Py_ssize_t *pos,
+    PyObject **key,
+    Py_hash_t *hash);
+
 // Export for '_pickle' shared extension
 PyAPI_FUNC(int) _PySet_Update(PyObject *set, PyObject *iterable);
 

Modules/_abc.c

@@ -862,15 +862,15 @@ subclasscheck_check_registry(_abc_data *impl, PyObject *subclass,
 
     // Make a local copy of the registry to protect against concurrent
     // modifications of _abc_registry.
-    PyObject *registry = PySet_New(registry_shared);
+    PyObject *registry = PyFrozenSet_New(registry_shared);
     if (registry == NULL) {
         return -1;
     }
     PyObject *key;
     Py_ssize_t pos = 0;
     Py_hash_t hash;
 
-    while (_PySet_NextEntry(registry, &pos, &key, &hash)) {
+    while (_PyFrozenSet_NextEntry(registry, &pos, &key, &hash)) {
         PyObject *rkey;
         if (PyWeakref_GetRef(key, &rkey) < 0) {
             // Someone inject non-weakref type in the registry.

Objects/codeobject.c

@@ -2347,7 +2347,7 @@ _PyCode_ConstantKey(PyObject *op)
             return NULL;
 
         i = 0;
-        while (_PySet_NextEntry(op, &pos, &item, &hash)) {
+        while (_PyFrozenSet_NextEntry(op, &pos, &item, &hash)) {
             PyObject *item_key;
 
             item_key = _PyCode_ConstantKey(item);

Objects/dictobject.c

@@ -2979,7 +2979,15 @@ dict_set_fromkeys(PyInterpreterState *interp, PyDictObject *mp,
         return NULL;
     }
 
-    while (_PySet_NextEntry(iterable, &pos, &key, &hash)) {
+    int (*next_entry_ptr)(PyObject *, Py_ssize_t *, PyObject **, Py_hash_t *);
+    if (PyFrozenSet_CheckExact(iterable)) {
+        next_entry_ptr = &_PyFrozenSet_NextEntry;
+    }
+    else {
+        next_entry_ptr = &_PySet_NextEntry;
+    }
+
+    while ((*next_entry_ptr)(iterable, &pos, &key, &hash)) {
         if (insertdict(interp, mp, Py_NewRef(key), hash, Py_NewRef(value))) {
             Py_DECREF(mp);
             return NULL;

Objects/listobject.c

@@ -1282,12 +1282,19 @@ list_extend_set(PyListObject *self, PySetObject *other)
     if (list_resize(self, m + n) < 0) {
         return -1;
     }
+    int (*next_entry_ptr)(PyObject *, Py_ssize_t *, PyObject **, Py_hash_t *);
+    if (PyFrozenSet_CheckExact(other)) {
+        next_entry_ptr = &_PyFrozenSet_NextEntry;
+    }
+    else {
+        next_entry_ptr = &_PySet_NextEntry;
+    }
     /* populate the end of self with iterable's items */
     Py_ssize_t setpos = 0;
     Py_hash_t hash;
     PyObject *key;
     PyObject **dest = self->ob_item + m;
-    while (_PySet_NextEntry((PyObject *)other, &setpos, &key, &hash)) {
+    while ((*next_entry_ptr)((PyObject *)other, &setpos, &key, &hash)) {
         Py_INCREF(key);
         FT_ATOMIC_STORE_PTR_RELEASE(*dest, key);
         dest++;

Objects/setobject.c

@@ -473,7 +473,7 @@ set_clear_internal(PySetObject *so)
  *     while (set_next(yourset, &pos, &entry)) {
  *              Refer to borrowed reference in entry->key.
  *     }
- *
+ *f
  * CAUTION:  In general, it isn't safe to use set_next in a loop that
  * mutates the table.
  */
@@ -2661,21 +2661,41 @@ PySet_Add(PyObject *anyset, PyObject *key)
     return rv;
 }
 
-// TODO: Make thread-safe in free-threaded builds
 int
-_PySet_NextEntry(PyObject *set, Py_ssize_t *pos, PyObject **key, Py_hash_t *hash)
+_PySet_NextEntry_lock_held(PyObject *set, Py_ssize_t *pos, PyObject **key, Py_hash_t *hash)
 {
     setentry *entry;
+    int ret = set_next((PySetObject *)set, pos, &entry);
+    if (ret == 0) {
+        return 0;
+    }
+    *key = entry->key;
+    *hash = entry->hash;
+    return 1;
+}
 
+int
+_PySet_NextEntry(PyObject *set, Py_ssize_t *pos, PyObject **key, Py_hash_t *hash)
+{
     if (!PyAnySet_Check(set)) {
         PyErr_BadInternalCall();
         return -1;
     }
-    if (set_next((PySetObject *)set, pos, &entry) == 0)
-        return 0;
-    *key = entry->key;
-    *hash = entry->hash;
-    return 1;
+    int ret;
+    Py_BEGIN_CRITICAL_SECTION(set);
+    ret = _PySet_NextEntry_lock_held(set, pos, key, hash);
+    Py_END_CRITICAL_SECTION();
+    return ret;
+}
+
+int
+_PyFrozenSet_NextEntry(PyObject *set, Py_ssize_t *pos, PyObject **key, Py_hash_t *hash)
+{
+    if (!PyFrozenSet_CheckExact(set)) {
+        PyErr_BadInternalCall();
+        return -1;
+    }
+    return _PySet_NextEntry_lock_held(set, pos, key, hash);
 }
 
 PyObject *

Python/compile.c

@@ -920,7 +920,7 @@ merge_consts_recursive(PyObject *const_cache, PyObject *o)
         Py_ssize_t i = 0, pos = 0;
         PyObject *item;
         Py_hash_t hash;
-        while (_PySet_NextEntry(o, &pos, &item, &hash)) {
+        while (_PyFrozenSet_NextEntry(o, &pos, &item, &hash)) {
             PyObject *k = merge_consts_recursive(const_cache, item);
             if (k == NULL) {
                 Py_DECREF(tuple);

Python/pylifecycle.c

@@ -2910,7 +2910,8 @@ _Py_DumpExtensionModules(int fd, PyInterpreterState *interp)
             Py_ssize_t i = 0;
             PyObject *item;
             Py_hash_t hash;
-            while (_PySet_NextEntry(stdlib_module_names, &i, &item, &hash)) {
+            // if stdlib_module_names is not NULL, it is always a frozenset.
+            while (_PyFrozenSet_NextEntry(stdlib_module_names, &i, &item, &hash)) {
                 if (PyUnicode_Check(item)
                     && PyUnicode_Compare(key, item) == 0)
                 {