diff --git a/Include/internal/pycore_critical_section.h b/Include/internal/pycore_critical_section.h index 7bebcdb67c7169..c6bbfdb108c469 100644 --- 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 index 5cb54cff9b76fd..c28efba4fef21a 100644 --- 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 new file mode 100644 index 00000000000000..9b0d8676c6bfa8 --- /dev/null +++ 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 index 556666779d8522..60cfae1099bb93 100644 --- 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);