python · eendebakpt · Jun 16, 2024 · Jun 16, 2024 · Jun 16, 2024 · Jun 16, 2024
diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h
@@ -98,6 +98,8 @@ extern "C" {
 # define Py_END_CRITICAL_SECTION()                                      \
         _PyCriticalSection_End(&_cs);                                   \
     }
+# define Py_EXIT_CRITICAL_SECTION()                                     \
+        _PyCriticalSection_End(&_cs);
 
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)                               \
     {                                                                   \
@@ -155,6 +157,7 @@ extern "C" {
 # define Py_BEGIN_CRITICAL_SECTION_MUT(mut)
 # define Py_BEGIN_CRITICAL_SECTION(op)
 # define Py_END_CRITICAL_SECTION()
+# define Py_EXIT_CRITICAL_SECTION()
 # define Py_BEGIN_CRITICAL_SECTION2(a, b)
 # define Py_END_CRITICAL_SECTION2()
 # define Py_BEGIN_CRITICAL_SECTION_SEQUENCE_FAST(original)

diff --git a/Lib/test/test_enumerate.py b/Lib/test/test_enumerate.py
@@ -3,8 +3,10 @@
 import sys
 import pickle
 import gc
+from threading import Thread
 
 from test import support
+from test.support import threading_helper
 
 class G:
     'Sequence using __getitem__'
@@ -292,5 +294,31 @@ def enum(self, iterable, start=sys.maxsize + 1):
                        (sys.maxsize+3,'c')]
 
 
+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
+        n = 100
+        start = sys.maxsize-10
+        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]))
+        _ = [t.start() for t in worker_threads]
+        _ = [t.join() for t in worker_threads]
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2024-06-16-19-21-07.gh-issue-120496.oP4xPi.rst b/Misc/NEWS.d/next/Core and Builtins/2024-06-16-19-21-07.gh-issue-120496.oP4xPi.rst
@@ -0,0 +1 @@
+Make :class:`enumerate` thread-safe.
diff --git a/Objects/enumobject.c b/Objects/enumobject.c
@@ -2,6 +2,7 @@
 
 #include "Python.h"
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
+#include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
 #include "pycore_long.h"          // _PyLong_GetOne()
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_object.h"        // _PyObject_GC_TRACK()
@@ -228,23 +229,38 @@ enum_next(enumobject *en)
     PyObject *it = en->en_sit;
     PyObject *old_index;
     PyObject *old_item;
+    int reuse_result = 0;
+
+    Py_BEGIN_CRITICAL_SECTION(en);
 
     next_item = (*Py_TYPE(it)->tp_iternext)(it);
-    if (next_item == NULL)
+    if (next_item == NULL) {
+        Py_EXIT_CRITICAL_SECTION();
         return NULL;
-
-    if (en->en_index == PY_SSIZE_T_MAX)
-        return enum_next_long(en, next_item);
+    }
+    if (en->en_index == PY_SSIZE_T_MAX) {
+        // we hold on to the lock for the entire call to enum_next_long
+        result = enum_next_long(en, next_item);
+        Py_EXIT_CRITICAL_SECTION();
+        return result;
+    }
 
     next_index = PyLong_FromSsize_t(en->en_index);
     if (next_index == NULL) {
         Py_DECREF(next_item);
+        Py_EXIT_CRITICAL_SECTION();
         return NULL;
     }
     en->en_index++;
 
-    if (Py_REFCNT(result) == 1) {
+    reuse_result = Py_REFCNT(result) == 1;
+    if (reuse_result) {
+        // increment within the critical section
         Py_INCREF(result);
+    }
+    Py_END_CRITICAL_SECTION();
+
+    if (reuse_result) {
         old_index = PyTuple_GET_ITEM(result, 0);
         old_item = PyTuple_GET_ITEM(result, 1);
         PyTuple_SET_ITEM(result, 0, next_index);