diff --git a/Lib/test/test_free_threading/test_itertools_batched.py b/Lib/test/test_free_threading/test_itertools.py similarity index 53% rename from Lib/test/test_free_threading/test_itertools_batched.py rename to Lib/test/test_free_threading/test_itertools.py index a754b4f9ea9902..8360afbf78cadd 100644 --- a/Lib/test/test_free_threading/test_itertools_batched.py +++ b/Lib/test/test_free_threading/test_itertools.py @@ -1,15 +1,15 @@ import unittest from threading import Thread, Barrier -from itertools import batched +from itertools import batched, cycle from test.support import threading_helper threading_helper.requires_working_threading(module=True) -class EnumerateThreading(unittest.TestCase): +class ItertoolsThreading(unittest.TestCase): @threading_helper.reap_threads - def test_threading(self): + def test_batched(self): number_of_threads = 10 number_of_iterations = 20 barrier = Barrier(number_of_threads) @@ -34,5 +34,31 @@ def work(it): barrier.reset() + @threading_helper.reap_threads + def test_cycle(self): + number_of_threads = 6 + number_of_iterations = 10 + number_of_cycles = 400 + + barrier = Barrier(number_of_threads) + def work(it): + barrier.wait() + for _ in range(number_of_cycles): + _ = next(it) + + data = (1, 2, 3, 4) + for it in range(number_of_iterations): + cycle_iterator = cycle(data) + worker_threads = [] + for ii in range(number_of_threads): + worker_threads.append( + Thread(target=work, args=[cycle_iterator])) + + with threading_helper.start_threads(worker_threads): + pass + + barrier.reset() + + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Library/2025-03-13-20-48-58.gh-issue-123471.cM4w4f.rst b/Misc/NEWS.d/next/Library/2025-03-13-20-48-58.gh-issue-123471.cM4w4f.rst new file mode 100644 index 00000000000000..cfc783900de70f --- /dev/null +++ b/Misc/NEWS.d/next/Library/2025-03-13-20-48-58.gh-issue-123471.cM4w4f.rst @@ -0,0 +1 @@ +Make concurrent iterations over :class:`itertools.cycle` safe under free-threading. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 943c1e8607b38f..7cfe5b468fd963 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -1124,7 +1124,6 @@ typedef struct { PyObject *it; PyObject *saved; Py_ssize_t index; - int firstpass; } cycleobject; #define cycleobject_CAST(op) ((cycleobject *)(op)) @@ -1165,8 +1164,7 @@ itertools_cycle_impl(PyTypeObject *type, PyObject *iterable) } lz->it = it; lz->saved = saved; - lz->index = 0; - lz->firstpass = 0; + lz->index = -1; return (PyObject *)lz; } @@ -1199,11 +1197,11 @@ cycle_next(PyObject *op) cycleobject *lz = cycleobject_CAST(op); PyObject *item; - if (lz->it != NULL) { + Py_ssize_t index = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->index); + + if (index < 0) { item = PyIter_Next(lz->it); if (item != NULL) { - if (lz->firstpass) - return item; if (PyList_Append(lz->saved, item)) { Py_DECREF(item); return NULL; @@ -1213,14 +1211,20 @@ cycle_next(PyObject *op) /* Note: StopIteration is already cleared by PyIter_Next() */ if (PyErr_Occurred()) return NULL; + index = 0; + FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, 0); +#ifndef Py_GIL_DISABLED Py_CLEAR(lz->it); +#endif } if (PyList_GET_SIZE(lz->saved) == 0) return NULL; - item = PyList_GET_ITEM(lz->saved, lz->index); - lz->index++; - if (lz->index >= PyList_GET_SIZE(lz->saved)) - lz->index = 0; + item = PyList_GET_ITEM(lz->saved, index); + index++; + if (index >= PyList_GET_SIZE(lz->saved)) { + index = 0; + } + FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, index); return Py_NewRef(item); }