Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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()
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Make concurrent iterations over :class:`itertools.cycle` safe under free-threading.
24 changes: 14 additions & 10 deletions Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -1124,7 +1124,6 @@ typedef struct {
PyObject *it;
PyObject *saved;
Py_ssize_t index;
int firstpass;
} cycleobject;

#define cycleobject_CAST(op) ((cycleobject *)(op))
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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;
Expand All @@ -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);
}

Expand Down
Loading