gh-112069: Make sets thread-safe with the GIL disabled (#113800)

This makes nearly all the operations on set thread-safe in the free-threaded build, with the exception of `_PySet_NextEntry` and `setiter_iternext`.

Co-authored-by: Sam Gross <colesbury@gmail.com>
Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>

2 files changed, 498 insertions, 152 deletions

diff --git a/Objects/clinic/setobject.c.h b/Objects/clinic/setobject.c.h
index f3c9699..3853ce3 100644
--- a/Objects/clinic/setobject.c.h
+++ b/Objects/clinic/setobject.c.h
@@ -2,6 +2,7 @@
 preserve
 [clinic start generated code]*/
 
+#include "pycore_critical_section.h"// Py_BEGIN_CRITICAL_SECTION()
 #include "pycore_modsupport.h"    // _PyArg_CheckPositional()
 
 PyDoc_STRVAR(set_pop__doc__,
@@ -21,7 +22,13 @@ set_pop_impl(PySetObject *so);
 static PyObject *
 set_pop(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return set_pop_impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_pop_impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
 
 PyDoc_STRVAR(set_update__doc__,
@@ -74,7 +81,13 @@ set_copy_impl(PySetObject *so);
 static PyObject *
 set_copy(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return set_copy_impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_copy_impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
 
 PyDoc_STRVAR(frozenset_copy__doc__,
@@ -92,7 +105,13 @@ frozenset_copy_impl(PySetObject *so);
 static PyObject *
 frozenset_copy(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return frozenset_copy_impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = frozenset_copy_impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
 
 PyDoc_STRVAR(set_clear__doc__,
@@ -110,7 +129,13 @@ set_clear_impl(PySetObject *so);
 static PyObject *
 set_clear(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return set_clear_impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_clear_impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
 
 PyDoc_STRVAR(set_union__doc__,
@@ -227,6 +252,21 @@ PyDoc_STRVAR(set_isdisjoint__doc__,
 #define SET_ISDISJOINT_METHODDEF    \
     {"isdisjoint", (PyCFunction)set_isdisjoint, METH_O, set_isdisjoint__doc__},
 
+static PyObject *
+set_isdisjoint_impl(PySetObject *so, PyObject *other);
+
+static PyObject *
+set_isdisjoint(PySetObject *so, PyObject *other)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    return_value = set_isdisjoint_impl(so, other);
+    Py_END_CRITICAL_SECTION2();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_difference_update__doc__,
 "difference_update($self, /, *others)\n"
 "--\n"
@@ -315,6 +355,21 @@ PyDoc_STRVAR(set_symmetric_difference__doc__,
 #define SET_SYMMETRIC_DIFFERENCE_METHODDEF    \
     {"symmetric_difference", (PyCFunction)set_symmetric_difference, METH_O, set_symmetric_difference__doc__},
 
+static PyObject *
+set_symmetric_difference_impl(PySetObject *so, PyObject *other);
+
+static PyObject *
+set_symmetric_difference(PySetObject *so, PyObject *other)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    return_value = set_symmetric_difference_impl(so, other);
+    Py_END_CRITICAL_SECTION2();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_issubset__doc__,
 "issubset($self, other, /)\n"
 "--\n"
@@ -324,6 +379,21 @@ PyDoc_STRVAR(set_issubset__doc__,
 #define SET_ISSUBSET_METHODDEF    \
     {"issubset", (PyCFunction)set_issubset, METH_O, set_issubset__doc__},
 
+static PyObject *
+set_issubset_impl(PySetObject *so, PyObject *other);
+
+static PyObject *
+set_issubset(PySetObject *so, PyObject *other)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    return_value = set_issubset_impl(so, other);
+    Py_END_CRITICAL_SECTION2();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_issuperset__doc__,
 "issuperset($self, other, /)\n"
 "--\n"
@@ -333,6 +403,21 @@ PyDoc_STRVAR(set_issuperset__doc__,
 #define SET_ISSUPERSET_METHODDEF    \
     {"issuperset", (PyCFunction)set_issuperset, METH_O, set_issuperset__doc__},
 
+static PyObject *
+set_issuperset_impl(PySetObject *so, PyObject *other);
+
+static PyObject *
+set_issuperset(PySetObject *so, PyObject *other)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    return_value = set_issuperset_impl(so, other);
+    Py_END_CRITICAL_SECTION2();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_add__doc__,
 "add($self, object, /)\n"
 "--\n"
@@ -344,6 +429,21 @@ PyDoc_STRVAR(set_add__doc__,
 #define SET_ADD_METHODDEF    \
     {"add", (PyCFunction)set_add, METH_O, set_add__doc__},
 
+static PyObject *
+set_add_impl(PySetObject *so, PyObject *key);
+
+static PyObject *
+set_add(PySetObject *so, PyObject *key)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_add_impl(so, key);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set___contains____doc__,
 "__contains__($self, object, /)\n"
 "--\n"
@@ -353,6 +453,21 @@ PyDoc_STRVAR(set___contains____doc__,
 #define SET___CONTAINS___METHODDEF    \
     {"__contains__", (PyCFunction)set___contains__, METH_O|METH_COEXIST, set___contains____doc__},
 
+static PyObject *
+set___contains___impl(PySetObject *so, PyObject *key);
+
+static PyObject *
+set___contains__(PySetObject *so, PyObject *key)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set___contains___impl(so, key);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_remove__doc__,
 "remove($self, object, /)\n"
 "--\n"
@@ -364,6 +479,21 @@ PyDoc_STRVAR(set_remove__doc__,
 #define SET_REMOVE_METHODDEF    \
     {"remove", (PyCFunction)set_remove, METH_O, set_remove__doc__},
 
+static PyObject *
+set_remove_impl(PySetObject *so, PyObject *key);
+
+static PyObject *
+set_remove(PySetObject *so, PyObject *key)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_remove_impl(so, key);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set_discard__doc__,
 "discard($self, object, /)\n"
 "--\n"
@@ -376,6 +506,21 @@ PyDoc_STRVAR(set_discard__doc__,
 #define SET_DISCARD_METHODDEF    \
     {"discard", (PyCFunction)set_discard, METH_O, set_discard__doc__},
 
+static PyObject *
+set_discard_impl(PySetObject *so, PyObject *key);
+
+static PyObject *
+set_discard(PySetObject *so, PyObject *key)
+{
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set_discard_impl(so, key);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
+}
+
 PyDoc_STRVAR(set___reduce____doc__,
 "__reduce__($self, /)\n"
 "--\n"
@@ -391,7 +536,13 @@ set___reduce___impl(PySetObject *so);
 static PyObject *
 set___reduce__(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return set___reduce___impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set___reduce___impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
 
 PyDoc_STRVAR(set___sizeof____doc__,
@@ -409,6 +560,12 @@ set___sizeof___impl(PySetObject *so);
 static PyObject *
 set___sizeof__(PySetObject *so, PyObject *Py_UNUSED(ignored))
 {
-    return set___sizeof___impl(so);
+    PyObject *return_value = NULL;
+
+    Py_BEGIN_CRITICAL_SECTION(so);
+    return_value = set___sizeof___impl(so);
+    Py_END_CRITICAL_SECTION();
+
+    return return_value;
 }
-/*[clinic end generated code: output=34a30591148da884 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=de4ee725bd29f758 input=a9049054013a1b77]*/
diff --git a/Objects/setobject.c b/Objects/setobject.c
index b4d803c..592711f 100644
--- a/Objects/setobject.c
+++ b/Objects/setobject.c
@@ -37,6 +37,7 @@
 #include "pycore_dict.h"                // _PyDict_Contains_KnownHash()
 #include "pycore_modsupport.h"          // _PyArg_NoKwnames()
 #include "pycore_object.h"              // _PyObject_GC_UNTRACK()
+#include "pycore_pyatomic_ft_wrappers.h"  // FT_ATOMIC_LOAD_SSIZE_RELAXED()
 #include "pycore_pyerrors.h"            // _PyErr_SetKeyError()
 #include "pycore_setobject.h"           // _PySet_NextEntry() definition
 #include <stddef.h>                     // offsetof()
@@ -130,6 +131,8 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
     int probes;
     int cmp;
 
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+
     /* Pre-increment is necessary to prevent arbitrary code in the rich
        comparison from deallocating the key just before the insertion. */
     Py_INCREF(key);
@@ -523,7 +526,7 @@ set_dealloc(PySetObject *so)
 }
 
 static PyObject *
-set_repr(PySetObject *so)
+set_repr_lock_held(PySetObject *so)
 {
     PyObject *result=NULL, *keys, *listrepr, *tmp;
     int status = Py_ReprEnter((PyObject*)so);
@@ -567,14 +570,24 @@ done:
     return result;
 }
 
+static PyObject *
+set_repr(PySetObject *so)
+{
+    PyObject *result;
+    Py_BEGIN_CRITICAL_SECTION(so);
+    result = set_repr_lock_held(so);
+    Py_END_CRITICAL_SECTION();
+    return result;
+}
+
 static Py_ssize_t
-set_len(PyObject *so)
+set_len(PySetObject *so)
 {
-    return ((PySetObject *)so)->used;
+    return FT_ATOMIC_LOAD_SSIZE_RELAXED(so->used);
 }
 
 static int
-set_merge(PySetObject *so, PyObject *otherset)
+set_merge_lock_held(PySetObject *so, PyObject *otherset)
 {
     PySetObject *other;
     PyObject *key;
@@ -584,6 +597,8 @@ set_merge(PySetObject *so, PyObject *otherset)
 
     assert (PyAnySet_Check(so));
     assert (PyAnySet_Check(otherset));
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(otherset);
 
     other = (PySetObject*)otherset;
     if (other == so || other->used == 0)
@@ -645,6 +660,7 @@ set_merge(PySetObject *so, PyObject *otherset)
 }
 
 /*[clinic input]
+@critical_section
 set.pop
     so: setobject
 
@@ -655,7 +671,7 @@ Raises KeyError if the set is empty.
 
 static PyObject *
 set_pop_impl(PySetObject *so)
-/*[clinic end generated code: output=4d65180f1271871b input=4a3f5552e660a260]*/
+/*[clinic end generated code: output=4d65180f1271871b input=9296c84921125060]*/
 {
     /* Make sure the search finger is in bounds */
     setentry *entry = so->table + (so->finger & so->mask);
@@ -889,58 +905,60 @@ PyTypeObject PySetIter_Type = {
 static PyObject *
 set_iter(PySetObject *so)
 {
+    Py_ssize_t size = set_len(so);
     setiterobject *si = PyObject_GC_New(setiterobject, &PySetIter_Type);
     if (si == NULL)
         return NULL;
     si->si_set = (PySetObject*)Py_NewRef(so);
-    si->si_used = so->used;
+    si->si_used = size;
     si->si_pos = 0;
-    si->len = so->used;
+    si->len = size;
     _PyObject_GC_TRACK(si);
     return (PyObject *)si;
 }
 
 static int
-set_update_internal(PySetObject *so, PyObject *other)
+set_update_dict_lock_held(PySetObject *so, PyObject *other)
 {
-    PyObject *key, *it;
+    assert(PyDict_CheckExact(other));
 
-    if (PyAnySet_Check(other))
-        return set_merge(so, other);
-
-    if (PyDict_CheckExact(other)) {
-        PyObject *value;
-        Py_ssize_t pos = 0;
-        Py_hash_t hash;
-        Py_ssize_t dictsize = PyDict_GET_SIZE(other);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
 
-        /* Do one big resize at the start, rather than
-        * incrementally resizing as we insert new keys.  Expect
-        * that there will be no (or few) overlapping keys.
-        */
-        if (dictsize < 0)
+    /* Do one big resize at the start, rather than
+    * incrementally resizing as we insert new keys.  Expect
+    * that there will be no (or few) overlapping keys.
+    */
+    Py_ssize_t dictsize = PyDict_GET_SIZE(other);
+    if ((so->fill + dictsize)*5 >= so->mask*3) {
+        if (set_table_resize(so, (so->used + dictsize)*2) != 0) {
             return -1;
-        if ((so->fill + dictsize)*5 >= so->mask*3) {
-            if (set_table_resize(so, (so->used + dictsize)*2) != 0)
-                return -1;
         }
-        int err = 0;
-        Py_BEGIN_CRITICAL_SECTION(other);
-        while (_PyDict_Next(other, &pos, &key, &value, &hash)) {
-            if (set_add_entry(so, key, hash)) {
-                err = -1;
-                goto exit;
-            }
+    }
+
+    Py_ssize_t pos = 0;
+    PyObject *key;
+    PyObject *value;
+    Py_hash_t hash;
+    while (_PyDict_Next(other, &pos, &key, &value, &hash)) {
+        if (set_add_entry(so, key, hash)) {
+            return -1;
         }
-exit:
-        Py_END_CRITICAL_SECTION();
-        return err;
     }
+    return 0;
+}
 
-    it = PyObject_GetIter(other);
-    if (it == NULL)
+static int
+set_update_iterable_lock_held(PySetObject *so, PyObject *other)
+{
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+
+    PyObject *it = PyObject_GetIter(other);
+    if (it == NULL) {
         return -1;
+    }
 
+    PyObject *key;
     while ((key = PyIter_Next(it)) != NULL) {
         if (set_add_key(so, key)) {
             Py_DECREF(it);
@@ -955,6 +973,69 @@ exit:
     return 0;
 }
 
+static int
+set_update_lock_held(PySetObject *so, PyObject *other)
+{
+    if (PyAnySet_Check(other)) {
+        return set_merge_lock_held(so, other);
+    }
+    else if (PyDict_CheckExact(other)) {
+        return set_update_dict_lock_held(so, other);
+    }
+    return set_update_iterable_lock_held(so, other);
+}
+
+// set_update for a `so` that is only visible to the current thread
+static int
+set_update_local(PySetObject *so, PyObject *other)
+{
+    assert(Py_REFCNT(so) == 1);
+    if (PyAnySet_Check(other)) {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION(other);
+        rv = set_merge_lock_held(so, other);
+        Py_END_CRITICAL_SECTION();
+        return rv;
+    }
+    else if (PyDict_CheckExact(other)) {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION(other);
+        rv = set_update_dict_lock_held(so, other);
+        Py_END_CRITICAL_SECTION();
+        return rv;
+    }
+    return set_update_iterable_lock_held(so, other);
+}
+
+static int
+set_update_internal(PySetObject *so, PyObject *other)
+{
+    if (PyAnySet_Check(other)) {
+        if (Py_Is((PyObject *)so, other)) {
+            return 0;
+        }
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION2(so, other);
+        rv = set_merge_lock_held(so, other);
+        Py_END_CRITICAL_SECTION2();
+        return rv;
+    }
+    else if (PyDict_CheckExact(other)) {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION2(so, other);
+        rv = set_update_dict_lock_held(so, other);
+        Py_END_CRITICAL_SECTION2();
+        return rv;
+    }
+    else {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION(so);
+        rv = set_update_iterable_lock_held(so, other);
+        Py_END_CRITICAL_SECTION();
+        return rv;
+    }
+}
+
 /*[clinic input]
 set.update
     so: setobject
@@ -1003,7 +1084,7 @@ make_new_set(PyTypeObject *type, PyObject *iterable)
     so->weakreflist = NULL;
 
     if (iterable != NULL) {
-        if (set_update_internal(so, iterable)) {
+        if (set_update_local(so, iterable)) {
             Py_DECREF(so);
             return NULL;
         }
@@ -1126,6 +1207,7 @@ set_swap_bodies(PySetObject *a, PySetObject *b)
 }
 
 /*[clinic input]
+@critical_section
 set.copy
     so: setobject
 
@@ -1134,12 +1216,22 @@ Return a shallow copy of a set.
 
 static PyObject *
 set_copy_impl(PySetObject *so)
-/*[clinic end generated code: output=c9223a1e1cc6b041 input=2b80b288d47b8cf1]*/
+/*[clinic end generated code: output=c9223a1e1cc6b041 input=c169a4fbb8209257]*/
 {
-    return make_new_set_basetype(Py_TYPE(so), (PyObject *)so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    PyObject *copy = make_new_set_basetype(Py_TYPE(so), NULL);
+    if (copy == NULL) {
+        return NULL;
+    }
+    if (set_merge_lock_held((PySetObject *)copy, (PyObject *)so) < 0) {
+        Py_DECREF(copy);
+        return NULL;
+    }
+    return copy;
 }
 
 /*[clinic input]
+@critical_section
 frozenset.copy
     so: setobject
 
@@ -1148,15 +1240,16 @@ Return a shallow copy of a set.
 
 static PyObject *
 frozenset_copy_impl(PySetObject *so)
-/*[clinic end generated code: output=b356263526af9e70 input=3dc65577d344eff7]*/
+/*[clinic end generated code: output=b356263526af9e70 input=fbf5bef131268dd7]*/
 {
     if (PyFrozenSet_CheckExact(so)) {
         return Py_NewRef(so);
     }
-    return set_copy(so, NULL);
+    return set_copy_impl(so);
 }
 
 /*[clinic input]
+@critical_section
 set.clear
     so: setobject
 
@@ -1165,7 +1258,7 @@ Remove all elements from this set.
 
 static PyObject *
 set_clear_impl(PySetObject *so)
-/*[clinic end generated code: output=4e71d5a83904161a input=74ac19794da81a39]*/
+/*[clinic end generated code: output=4e71d5a83904161a input=c6f831b366111950]*/
 {
     set_clear_internal(so);
     Py_RETURN_NONE;
@@ -1196,7 +1289,7 @@ set_union_impl(PySetObject *so, PyObject *args)
         other = PyTuple_GET_ITEM(args, i);
         if ((PyObject *)so == other)
             continue;
-        if (set_update_internal(result, other)) {
+        if (set_update_local(result, other)) {
             Py_DECREF(result);
             return NULL;
         }
@@ -1213,11 +1306,13 @@ set_or(PySetObject *so, PyObject *other)
         Py_RETURN_NOTIMPLEMENTED;
 
     result = (PySetObject *)set_copy(so, NULL);
-    if (result == NULL)
+    if (result == NULL) {
         return NULL;
-    if ((PyObject *)so == other)
+    }
+    if (Py_Is((PyObject *)so, other)) {
         return (PyObject *)result;
-    if (set_update_internal(result, other)) {
+    }
+    if (set_update_local(result, other)) {
         Py_DECREF(result);
         return NULL;
     }
@@ -1230,8 +1325,9 @@ set_ior(PySetObject *so, PyObject *other)
     if (!PyAnySet_Check(other))
         Py_RETURN_NOTIMPLEMENTED;
 
-    if (set_update_internal(so, other))
+    if (set_update_internal(so, other)) {
         return NULL;
+    }
     return Py_NewRef(so);
 }
 
@@ -1244,7 +1340,7 @@ set_intersection(PySetObject *so, PyObject *other)
     int rv;
 
     if ((PyObject *)so == other)
-        return set_copy(so, NULL);
+        return set_copy_impl(so);
 
     result = (PySetObject *)make_new_set_basetype(Py_TYPE(so), NULL);
     if (result == NULL)
@@ -1333,13 +1429,17 @@ set_intersection_multi_impl(PySetObject *so, PyObject *args)
 {
     Py_ssize_t i;
 
-    if (PyTuple_GET_SIZE(args) == 0)
+    if (PyTuple_GET_SIZE(args) == 0) {
         return set_copy(so, NULL);
+    }
 
     PyObject *result = Py_NewRef(so);
     for (i=0 ; i<PyTuple_GET_SIZE(args) ; i++) {
         PyObject *other = PyTuple_GET_ITEM(args, i);
-        PyObject *newresult = set_intersection((PySetObject *)result, other);
+        PyObject *newresult;
+        Py_BEGIN_CRITICAL_SECTION2(result, other);
+        newresult = set_intersection((PySetObject *)result, other);
+        Py_END_CRITICAL_SECTION2();
         if (newresult == NULL) {
             Py_DECREF(result);
             return NULL;
@@ -1380,7 +1480,9 @@ set_intersection_update_multi_impl(PySetObject *so, PyObject *args)
     tmp = set_intersection_multi_impl(so, args);
     if (tmp == NULL)
         return NULL;
+    Py_BEGIN_CRITICAL_SECTION(so);
     set_swap_bodies(so, (PySetObject *)tmp);
+    Py_END_CRITICAL_SECTION();
     Py_DECREF(tmp);
     Py_RETURN_NONE;
 }
@@ -1390,7 +1492,13 @@ set_and(PySetObject *so, PyObject *other)
 {
     if (!PyAnySet_Check(so) || !PyAnySet_Check(other))
         Py_RETURN_NOTIMPLEMENTED;
-    return set_intersection(so, other);
+
+    PyObject *rv;
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    rv = set_intersection(so, other);
+    Py_END_CRITICAL_SECTION2();
+
+    return rv;
 }
 
 static PyObject *
@@ -1400,7 +1508,11 @@ set_iand(PySetObject *so, PyObject *other)
 
     if (!PyAnySet_Check(other))
         Py_RETURN_NOTIMPLEMENTED;
+
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
     result = set_intersection_update(so, other);
+    Py_END_CRITICAL_SECTION2();
+
     if (result == NULL)
         return NULL;
     Py_DECREF(result);
@@ -1408,6 +1520,7 @@ set_iand(PySetObject *so, PyObject *other)
 }
 
 /*[clinic input]
+@critical_section so other
 set.isdisjoint
     so: setobject
     other: object
@@ -1417,8 +1530,8 @@ Return True if two sets have a null intersection.
 [clinic start generated code]*/
 
 static PyObject *
-set_isdisjoint(PySetObject *so, PyObject *other)
-/*[clinic end generated code: output=a92bbf9a2db6a3da input=c254ddec8a2326e3]*/
+set_isdisjoint_impl(PySetObject *so, PyObject *other)
+/*[clinic end generated code: output=273493f2d57c565e input=32f8dcab5e0fc7d6]*/
 {
     PyObject *key, *it, *tmp;
     int rv;
@@ -1479,6 +1592,9 @@ set_isdisjoint(PySetObject *so, PyObject *other)
 static int
 set_difference_update_internal(PySetObject *so, PyObject *other)
 {
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+
     if ((PyObject *)so == other)
         return set_clear_internal(so);
 
@@ -1551,8 +1667,13 @@ set_difference_update_impl(PySetObject *so, PyObject *args)
 
     for (i=0 ; i<PyTuple_GET_SIZE(args) ; i++) {
         PyObject *other = PyTuple_GET_ITEM(args, i);
-        if (set_difference_update_internal(so, other))
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION2(so, other);
+        rv = set_difference_update_internal(so, other);
+        Py_END_CRITICAL_SECTION2();
+        if (rv) {
             return NULL;
+        }
     }
     Py_RETURN_NONE;
 }
@@ -1562,7 +1683,7 @@ set_copy_and_difference(PySetObject *so, PyObject *other)
 {
     PyObject *result;
 
-    result = set_copy(so, NULL);
+    result = set_copy_impl(so);
     if (result == NULL)
         return NULL;
     if (set_difference_update_internal((PySetObject *) result, other) == 0)
@@ -1663,17 +1784,24 @@ set_difference_multi_impl(PySetObject *so, PyObject *args)
     Py_ssize_t i;
     PyObject *result, *other;
 
-    if (PyTuple_GET_SIZE(args) == 0)
+    if (PyTuple_GET_SIZE(args) == 0) {
         return set_copy(so, NULL);
+    }
 
     other = PyTuple_GET_ITEM(args, 0);
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
     result = set_difference(so, other);
+    Py_END_CRITICAL_SECTION2();
     if (result == NULL)
         return NULL;
 
     for (i=1 ; i<PyTuple_GET_SIZE(args) ; i++) {
         other = PyTuple_GET_ITEM(args, i);
-        if (set_difference_update_internal((PySetObject *)result, other)) {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION(other);
+        rv = set_difference_update_internal((PySetObject *)result, other);
+        Py_END_CRITICAL_SECTION();
+        if (rv) {
             Py_DECREF(result);
             return NULL;
         }
@@ -1686,7 +1814,12 @@ set_sub(PySetObject *so, PyObject *other)
 {
     if (!PyAnySet_Check(so) || !PyAnySet_Check(other))
         Py_RETURN_NOTIMPLEMENTED;
-    return set_difference(so, other);
+
+    PyObject *rv;
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    rv = set_difference(so, other);
+    Py_END_CRITICAL_SECTION2();
+    return rv;
 }
 
 static PyObject *
@@ -1694,36 +1827,69 @@ set_isub(PySetObject *so, PyObject *other)
 {
     if (!PyAnySet_Check(other))
         Py_RETURN_NOTIMPLEMENTED;
-    if (set_difference_update_internal(so, other))
+
+    int rv;
+    Py_BEGIN_CRITICAL_SECTION2(so, other);
+    rv = set_difference_update_internal(so, other);
+    Py_END_CRITICAL_SECTION2();
+    if (rv < 0) {
         return NULL;
+    }
     return Py_NewRef(so);
 }
 
-static PyObject *
+static int
 set_symmetric_difference_update_dict(PySetObject *so, PyObject *other)
 {
-    PyObject *key;
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+
     Py_ssize_t pos = 0;
+    PyObject *key, *value;
     Py_hash_t hash;
-    PyObject *value;
-    int rv;
-
     while (_PyDict_Next(other, &pos, &key, &value, &hash)) {
         Py_INCREF(key);
-        rv = set_discard_entry(so, key, hash);
+        int rv = set_discard_entry(so, key, hash);
         if (rv < 0) {
             Py_DECREF(key);
-            return NULL;
+            return -1;
         }
         if (rv == DISCARD_NOTFOUND) {
             if (set_add_entry(so, key, hash)) {
                 Py_DECREF(key);
-                return NULL;
+                return -1;
             }
         }
         Py_DECREF(key);
     }
-    Py_RETURN_NONE;
+    return 0;
+}
+
+static int
+set_symmetric_difference_update_set(PySetObject *so, PySetObject *other)
+{
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+
+    Py_ssize_t pos = 0;
+    setentry *entry;
+    while (set_next(other, &pos, &entry)) {
+        PyObject *key = Py_NewRef(entry->key);
+        Py_hash_t hash = entry->hash;
+        int rv = set_discard_entry(so, key, hash);
+        if (rv < 0) {
+            Py_DECREF(key);
+            return -1;
+        }
+        if (rv == DISCARD_NOTFOUND) {
+            if (set_add_entry(so, key, hash)) {
+                Py_DECREF(key);
+                return -1;
+            }
+        }
+        Py_DECREF(key);
+    }
+    return 0;
 }
 
 /*[clinic input]
@@ -1739,58 +1905,41 @@ static PyObject *
 set_symmetric_difference_update(PySetObject *so, PyObject *other)
 /*[clinic end generated code: output=fbb049c0806028de input=a50acf0365e1f0a5]*/
 {
-    PySetObject *otherset;
-    PyObject *key;
-    Py_ssize_t pos = 0;
-    Py_hash_t hash;
-    setentry *entry;
-    int rv;
-
-    if ((PyObject *)so == other)
+    if (Py_Is((PyObject *)so, other)) {
         return set_clear(so, NULL);
+    }
 
+    int rv;
     if (PyDict_CheckExact(other)) {
-        PyObject *res;
-
-        Py_BEGIN_CRITICAL_SECTION(other);
-        res = set_symmetric_difference_update_dict(so, other);
-        Py_END_CRITICAL_SECTION();
-
-        return res;
+        Py_BEGIN_CRITICAL_SECTION2(so, other);
+        rv = set_symmetric_difference_update_dict(so, other);
+        Py_END_CRITICAL_SECTION2();
     }
-
-    if (PyAnySet_Check(other)) {
-        otherset = (PySetObject *)Py_NewRef(other);
-    } else {
-        otherset = (PySetObject *)make_new_set_basetype(Py_TYPE(so), other);
-        if (otherset == NULL)
-            return NULL;
+    else if (PyAnySet_Check(other)) {
+        Py_BEGIN_CRITICAL_SECTION2(so, other);
+        rv = set_symmetric_difference_update_set(so, (PySetObject *)other);
+        Py_END_CRITICAL_SECTION2();
     }
-
-    while (set_next(otherset, &pos, &entry)) {
-        key = entry->key;
-        hash = entry->hash;
-        Py_INCREF(key);
-        rv = set_discard_entry(so, key, hash);
-        if (rv < 0) {
-            Py_DECREF(otherset);
-            Py_DECREF(key);
+    else {
+        PySetObject *otherset = (PySetObject *)make_new_set_basetype(Py_TYPE(so), other);
+        if (otherset == NULL) {
             return NULL;
         }
-        if (rv == DISCARD_NOTFOUND) {
-            if (set_add_entry(so, key, hash)) {
-                Py_DECREF(otherset);
-                Py_DECREF(key);
-                return NULL;
-            }
-        }
-        Py_DECREF(key);
+
+        Py_BEGIN_CRITICAL_SECTION(so);
+        rv = set_symmetric_difference_update_set(so, otherset);
+        Py_END_CRITICAL_SECTION();
+
+        Py_DECREF(otherset);
+    }
+    if (rv < 0) {
+        return NULL;
     }
-    Py_DECREF(otherset);
     Py_RETURN_NONE;
 }
 
 /*[clinic input]
+@critical_section so other
 set.symmetric_difference
     so: setobject
     other: object
@@ -1800,22 +1949,22 @@ Return a new set with elements in either the set or other but not both.
 [clinic start generated code]*/
 
 static PyObject *
-set_symmetric_difference(PySetObject *so, PyObject *other)
-/*[clinic end generated code: output=f95364211b88775a input=f18af370ad72ebac]*/
+set_symmetric_difference_impl(PySetObject *so, PyObject *other)
+/*[clinic end generated code: output=270ee0b5d42b0797 input=624f6e7bbdf70db1]*/
 {
-    PyObject *rv;
-    PySetObject *otherset;
-
-    otherset = (PySetObject *)make_new_set_basetype(Py_TYPE(so), other);
-    if (otherset == NULL)
+    PySetObject *result = (PySetObject *)make_new_set_basetype(Py_TYPE(so), NULL);
+    if (result == NULL) {
         return NULL;
-    rv = set_symmetric_difference_update(otherset, (PyObject *)so);
-    if (rv == NULL) {
-        Py_DECREF(otherset);
+    }
+    if (set_update_lock_held(result, other) < 0) {
+        Py_DECREF(result);
         return NULL;
     }
-    Py_DECREF(rv);
-    return (PyObject *)otherset;
+    if (set_symmetric_difference_update_set(result, so) < 0) {
+        Py_DECREF(result);
+        return NULL;
+    }
+    return (PyObject *)result;
 }
 
 static PyObject *
@@ -1841,6 +1990,7 @@ set_ixor(PySetObject *so, PyObject *other)
 }
 
 /*[clinic input]
+@critical_section so other
 set.issubset
     so: setobject
     other: object
@@ -1850,8 +2000,8 @@ Report whether another set contains this set.
 [clinic start generated code]*/
 
 static PyObject *
-set_issubset(PySetObject *so, PyObject *other)
-/*[clinic end generated code: output=78aef1f377aedef1 input=37fbc579b609db0c]*/
+set_issubset_impl(PySetObject *so, PyObject *other)
+/*[clinic end generated code: output=b2b59d5f314555ce input=f2a4fd0f2537758b]*/
 {
     setentry *entry;
     Py_ssize_t pos = 0;
@@ -1885,6 +2035,7 @@ set_issubset(PySetObject *so, PyObject *other)
 }
 
 /*[clinic input]
+@critical_section so other
 set.issuperset
     so: setobject
     other: object
@@ -1894,8 +2045,8 @@ Report whether this set contains another set.
 [clinic start generated code]*/
 
 static PyObject *
-set_issuperset(PySetObject *so, PyObject *other)
-/*[clinic end generated code: output=7d2b71dd714a7ec7 input=fd5dab052f2e9bb3]*/
+set_issuperset_impl(PySetObject *so, PyObject *other)
+/*[clinic end generated code: output=ecf00ce552c09461 input=5f2e1f262e6e4ccc]*/
 {
     if (PyAnySet_Check(other)) {
         return set_issubset((PySetObject *)other, (PyObject *)so);
@@ -1924,6 +2075,7 @@ set_issuperset(PySetObject *so, PyObject *other)
     Py_RETURN_TRUE;
 }
 
+// TODO: Make thread-safe in free-threaded builds
 static PyObject *
 set_richcompare(PySetObject *v, PyObject *w, int op)
 {
@@ -1968,6 +2120,7 @@ set_richcompare(PySetObject *v, PyObject *w, int op)
 }
 
 /*[clinic input]
+@critical_section
 set.add
     so: setobject
     object as key: object
@@ -1979,16 +2132,16 @@ This has no effect if the element is already present.
 [clinic start generated code]*/
 
 static PyObject *
-set_add(PySetObject *so, PyObject *key)
-/*[clinic end generated code: output=cd9c2d5c2069c2ba input=96f1efe029e47972]*/
+set_add_impl(PySetObject *so, PyObject *key)
+/*[clinic end generated code: output=4cc4a937f1425c96 input=03baf62cb0e66514]*/
 {
     if (set_add_key(so, key))
         return NULL;
     Py_RETURN_NONE;
 }
 
-int
-_PySet_Contains(PySetObject *so, PyObject *key)
+static int
+set_contains_lock_held(PySetObject *so, PyObject *key)
 {
     PyObject *tmpkey;
     int rv;
@@ -2007,7 +2160,18 @@ _PySet_Contains(PySetObject *so, PyObject *key)
     return rv;
 }
 
+int
+_PySet_Contains(PySetObject *so, PyObject *key)
+{
+    int rv;
+    Py_BEGIN_CRITICAL_SECTION(so);
+    rv = set_contains_lock_held(so, key);
+    Py_END_CRITICAL_SECTION();
+    return rv;
+}
+
 /*[clinic input]
+@critical_section
 @coexist
 set.__contains__
     so: setobject
@@ -2018,18 +2182,19 @@ x.__contains__(y) <==> y in x.
 [clinic start generated code]*/
 
 static PyObject *
-set___contains__(PySetObject *so, PyObject *key)
-/*[clinic end generated code: output=b5948bc5c590d3ca input=cf4c72db704e4cf0]*/
+set___contains___impl(PySetObject *so, PyObject *key)
+/*[clinic end generated code: output=b44863d034b3c70e input=4a7d568459617f24]*/
 {
     long result;
 
-    result = _PySet_Contains(so, key);
+    result = set_contains_lock_held(so, key);
     if (result < 0)
         return NULL;
     return PyBool_FromLong(result);
 }
 
 /*[clinic input]
+@critical_section
 set.remove
     so: setobject
     object as key: object
@@ -2041,8 +2206,8 @@ If the element is not a member, raise a KeyError.
 [clinic start generated code]*/
 
 static PyObject *
-set_remove(PySetObject *so, PyObject *key)
-/*[clinic end generated code: output=08ae496d0cd2b8c1 input=10132515dfe8ebd7]*/
+set_remove_impl(PySetObject *so, PyObject *key)
+/*[clinic end generated code: output=0b9134a2a2200363 input=893e1cb1df98227a]*/
 {
     PyObject *tmpkey;
     int rv;
@@ -2069,6 +2234,7 @@ set_remove(PySetObject *so, PyObject *key)
 }
 
 /*[clinic input]
+@critical_section
 set.discard
     so: setobject
     object as key: object
@@ -2081,8 +2247,8 @@ an exception when an element is missing from the set.
 [clinic start generated code]*/
 
 static PyObject *
-set_discard(PySetObject *so, PyObject *key)
-/*[clinic end generated code: output=9181b60d7bb7d480 input=82a689eba94d5ad9]*/
+set_discard_impl(PySetObject *so, PyObject *key)
+/*[clinic end generated code: output=eec3b687bf32759e input=861cb7fb69b4def0]*/
 {
     PyObject *tmpkey;
     int rv;
@@ -2104,6 +2270,7 @@ set_discard(PySetObject *so, PyObject *key)
 }
 
 /*[clinic input]
+@critical_section
 set.__reduce__
     so: setobject
 
@@ -2112,7 +2279,7 @@ Return state information for pickling.
 
 static PyObject *
 set___reduce___impl(PySetObject *so)
-/*[clinic end generated code: output=9af7d0e029df87ee input=531375e87a24a449]*/
+/*[clinic end generated code: output=9af7d0e029df87ee input=59405a4249e82f71]*/
 {
     PyObject *keys=NULL, *args=NULL, *result=NULL, *state=NULL;
 
@@ -2134,6 +2301,7 @@ done:
 }
 
 /*[clinic input]
+@critical_section
 set.__sizeof__
     so: setobject
 
@@ -2142,7 +2310,7 @@ S.__sizeof__() -> size of S in memory, in bytes.
 
 static PyObject *
 set___sizeof___impl(PySetObject *so)
-/*[clinic end generated code: output=4bfa3df7bd38ed88 input=0f214fc2225319fc]*/
+/*[clinic end generated code: output=4bfa3df7bd38ed88 input=09e1a09f168eaa23]*/
 {
     size_t res = _PyObject_SIZE(Py_TYPE(so));
     if (so->table != so->smalltable) {
@@ -2156,13 +2324,17 @@ set_init(PySetObject *self, PyObject *args, PyObject *kwds)
 {
     PyObject *iterable = NULL;
 
-     if (!_PyArg_NoKeywords("set", kwds))
+    if (!_PyArg_NoKeywords("set", kwds))
         return -1;
     if (!PyArg_UnpackTuple(args, Py_TYPE(self)->tp_name, 0, 1, &iterable))
         return -1;
+
+    Py_BEGIN_CRITICAL_SECTION(self);
     if (self->fill)
         set_clear_internal(self);
     self->hash = -1;
+    Py_END_CRITICAL_SECTION();
+
     if (iterable == NULL)
         return 0;
     return set_update_internal(self, iterable);
@@ -2191,7 +2363,7 @@ set_vectorcall(PyObject *type, PyObject * const*args,
 }
 
 static PySequenceMethods set_as_sequence = {
-    set_len,                            /* sq_length */
+    (lenfunc)set_len,                   /* sq_length */
     0,                                  /* sq_concat */
     0,                                  /* sq_repeat */
     0,                                  /* sq_item */
@@ -2424,7 +2596,7 @@ PySet_Size(PyObject *anyset)
         PyErr_BadInternalCall();
         return -1;
     }
-    return PySet_GET_SIZE(anyset);
+    return set_len((PySetObject *)anyset);
 }
 
 int
@@ -2434,7 +2606,8 @@ PySet_Clear(PyObject *set)
         PyErr_BadInternalCall();
         return -1;
     }
-    return set_clear_internal((PySetObject *)set);
+    (void)set_clear((PySetObject *)set, NULL);
+    return 0;
 }
 
 int
@@ -2444,7 +2617,12 @@ PySet_Contains(PyObject *anyset, PyObject *key)
         PyErr_BadInternalCall();
         return -1;
     }
-    return set_contains_key((PySetObject *)anyset, key);
+
+    int rv;
+    Py_BEGIN_CRITICAL_SECTION(anyset);
+    rv = set_contains_key((PySetObject *)anyset, key);
+    Py_END_CRITICAL_SECTION();
+    return rv;
 }
 
 int
@@ -2454,7 +2632,12 @@ PySet_Discard(PyObject *set, PyObject *key)
         PyErr_BadInternalCall();
         return -1;
     }
-    return set_discard_key((PySetObject *)set, key);
+
+    int rv;
+    Py_BEGIN_CRITICAL_SECTION(set);
+    rv = set_discard_key((PySetObject *)set, key);
+    Py_END_CRITICAL_SECTION();
+    return rv;
 }
 
 int
@@ -2465,9 +2648,15 @@ PySet_Add(PyObject *anyset, PyObject *key)
         PyErr_BadInternalCall();
         return -1;
     }
-    return set_add_key((PySetObject *)anyset, key);
+
+    int rv;
+    Py_BEGIN_CRITICAL_SECTION(anyset);
+    rv = set_add_key((PySetObject *)anyset, key);
+    Py_END_CRITICAL_SECTION();
+    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)
 {