gh-120496: Make enum_iter thread safe

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)

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()

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.

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);