gh-123471: make concurrent iteration over itertools.cycle safe under free-threading (#131212)

Co-authored-by: Kumar Aditya <kumaraditya@python.org>
This commit is contained in:
Pieter Eendebak 2025-06-02 16:43:32 +02:00 committed by GitHub
parent b6237c3602
commit 26a1cd4e8c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 46 additions and 14 deletions

View file

@ -1,15 +1,15 @@
import unittest import unittest
from threading import Thread, Barrier from threading import Thread, Barrier
from itertools import batched from itertools import batched, cycle
from test.support import threading_helper from test.support import threading_helper
threading_helper.requires_working_threading(module=True) threading_helper.requires_working_threading(module=True)
class EnumerateThreading(unittest.TestCase): class ItertoolsThreading(unittest.TestCase):
@threading_helper.reap_threads @threading_helper.reap_threads
def test_threading(self): def test_batched(self):
number_of_threads = 10 number_of_threads = 10
number_of_iterations = 20 number_of_iterations = 20
barrier = Barrier(number_of_threads) barrier = Barrier(number_of_threads)
@ -34,5 +34,31 @@ class EnumerateThreading(unittest.TestCase):
barrier.reset() 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__": if __name__ == "__main__":
unittest.main() unittest.main()

View file

@ -0,0 +1 @@
Make concurrent iterations over :class:`itertools.cycle` safe under free-threading.

View file

@ -1124,7 +1124,6 @@ typedef struct {
PyObject *it; PyObject *it;
PyObject *saved; PyObject *saved;
Py_ssize_t index; Py_ssize_t index;
int firstpass;
} cycleobject; } cycleobject;
#define cycleobject_CAST(op) ((cycleobject *)(op)) #define cycleobject_CAST(op) ((cycleobject *)(op))
@ -1165,8 +1164,7 @@ itertools_cycle_impl(PyTypeObject *type, PyObject *iterable)
} }
lz->it = it; lz->it = it;
lz->saved = saved; lz->saved = saved;
lz->index = 0; lz->index = -1;
lz->firstpass = 0;
return (PyObject *)lz; return (PyObject *)lz;
} }
@ -1199,11 +1197,11 @@ cycle_next(PyObject *op)
cycleobject *lz = cycleobject_CAST(op); cycleobject *lz = cycleobject_CAST(op);
PyObject *item; 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); item = PyIter_Next(lz->it);
if (item != NULL) { if (item != NULL) {
if (lz->firstpass)
return item;
if (PyList_Append(lz->saved, item)) { if (PyList_Append(lz->saved, item)) {
Py_DECREF(item); Py_DECREF(item);
return NULL; return NULL;
@ -1213,15 +1211,22 @@ cycle_next(PyObject *op)
/* Note: StopIteration is already cleared by PyIter_Next() */ /* Note: StopIteration is already cleared by PyIter_Next() */
if (PyErr_Occurred()) if (PyErr_Occurred())
return NULL; return NULL;
index = 0;
FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, 0);
#ifndef Py_GIL_DISABLED
Py_CLEAR(lz->it); Py_CLEAR(lz->it);
#endif
} }
if (PyList_GET_SIZE(lz->saved) == 0) if (PyList_GET_SIZE(lz->saved) == 0)
return NULL; return NULL;
item = PyList_GET_ITEM(lz->saved, lz->index); item = PyList_GetItemRef(lz->saved, index);
lz->index++; assert(item);
if (lz->index >= PyList_GET_SIZE(lz->saved)) index++;
lz->index = 0; if (index >= PyList_GET_SIZE(lz->saved)) {
return Py_NewRef(item); index = 0;
}
FT_ATOMIC_STORE_SSIZE_RELAXED(lz->index, index);
return item;
} }
static PyType_Slot cycle_slots[] = { static PyType_Slot cycle_slots[] = {