Make concurrent iteration over enumerate safe under free-threading

python · kumaraditya303 · Mar 13, 2025 · Oct 19, 2024 · Oct 19, 2024 · Oct 19, 2024
commit 609ce86a27056d6cf537e9fbde9ce90b0d5c0b6f
@@ -0,0 +1,41 @@
+import unittest
+import sys
+from threading import Thread
+
+from test.support import threading_helper
+
+
+class EnumerateThreading(unittest.TestCase):
+    @staticmethod
+    def work(enum, start):
+        while True:
+            try:
+                value = next(enum)
+            except StopIteration:
+                break
+            else:
+                if value[0] + start != value[1]:
+                    raise ValueError(f'enumerate returned pair {value}')
+
+    @threading_helper.reap_threads
+    @threading_helper.requires_working_threading()
+    def test_threading(self):
+        number_of_threads = 4
+        number_of_iterations = 8
+        n = 100
+        start = sys.maxsize - 40
+
+        for _ in range(number_of_iterations):
+            enum = enumerate(range(start, start + n))
+            worker_threads = []
+            for ii in range(number_of_threads):
+                worker_threads.append(
+                    Thread(target=self.work, args=[enum, start]))
+            for t in worker_threads:
+                t.start()
+            for t in worker_threads:
+                t.join()
+
+
+if __name__ == "__main__":
+    unittest.main()
@@ -177,6 +177,7 @@ enum_next_long(enumobject *en, PyObject* next_item)
     PyObject *old_index;
     PyObject *old_item;
 
+    Py_BEGIN_CRITICAL_SECTION(en);
     if (en->en_longindex == NULL) {
         en->en_longindex = PyLong_FromSsize_t(PY_SSIZE_T_MAX);
         if (en->en_longindex == NULL) {
@@ -192,8 +193,9 @@ enum_next_long(enumobject *en, PyObject* next_item)
         return NULL;
     }
     en->en_longindex = stepped_up;
+    Py_END_CRITICAL_SECTION();
 
-    if (Py_REFCNT(result) == 1) {
+    if (_PyObject_IsUniquelyReferenced(result)) {
         Py_INCREF(result);
         old_index = PyTuple_GET_ITEM(result, 0);
         old_item = PyTuple_GET_ITEM(result, 1);
@@ -233,17 +235,18 @@ enum_next(enumobject *en)
     if (next_item == NULL)
         return NULL;
 
-    if (en->en_index == PY_SSIZE_T_MAX)
+    Py_ssize_t en_index = FT_ATOMIC_LOAD_SSIZE_RELAXED(en->en_index);
+    if (en_index == PY_SSIZE_T_MAX)
         return enum_next_long(en, next_item);
 
-    next_index = PyLong_FromSsize_t(en->en_index);
+    next_index = PyLong_FromSsize_t(en_index);
     if (next_index == NULL) {
         Py_DECREF(next_item);
         return NULL;
     }
-    en->en_index++;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(en->en_index, en_index + 1);
 
-    if (Py_REFCNT(result) == 1) {
+    if (_PyObject_IsUniquelyReferenced(result)) {
         Py_INCREF(result);
         old_index = PyTuple_GET_ITEM(result, 0);
         old_item = PyTuple_GET_ITEM(result, 1);
@@ -272,10 +275,12 @@ enum_next(enumobject *en)
 static PyObject *
 enum_reduce(enumobject *en, PyObject *Py_UNUSED(ignored))
 {
+    Py_BEGIN_CRITICAL_SECTION(en);
     if (en->en_longindex != NULL)
         return Py_BuildValue("O(OO)", Py_TYPE(en), en->en_sit, en->en_longindex);
     else
         return Py_BuildValue("O(On)", Py_TYPE(en), en->en_sit, en->en_index);
+    Py_END_CRITICAL_SECTION();
 }
 
 PyDoc_STRVAR(reduce_doc, "Return state information for pickling.");