Skip to content

Commit 6c389c4

Browse files
gh-153062: Fix a crash iterating itertools.tee on the free-threaded build (gh-153063)
itertools.tee branches share a linked list of teedataobject cells. On the free-threaded build, iterating one branch from multiple threads, or iterating sibling branches concurrently, raced on the shared cells and each branch's position, corrupting refcounts and crashing. Lock each teedataobject while reading, extending, or clearing it, and snapshot each branch's position under the tee object's own lock, revalidating before advancing, so the two locks are never nested. Concurrent iteration of one tee is undefined and may raise RuntimeError as documented, but no longer crashes. The free-threading snapshot and revalidation add per-element overhead on the default build, where the GIL already serializes access. Guard that path under Py_GIL_DISABLED so the default build keeps the original iteration and the free-threaded path is unchanged. Co-authored-by: Neil Schemenauer <nas-github@arctrix.com>
1 parent d83d267 commit 6c389c4

3 files changed

Lines changed: 157 additions & 15 deletions

File tree

Lib/test/test_free_threading/test_itertools.py

Lines changed: 80 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,14 @@
11
import unittest
2-
from itertools import accumulate, batched, chain, combinations_with_replacement, cycle, permutations, zip_longest
2+
from itertools import (
3+
accumulate,
4+
batched,
5+
chain,
6+
combinations_with_replacement,
7+
cycle,
8+
permutations,
9+
tee,
10+
zip_longest,
11+
)
312
from test.support import threading_helper
413

514

@@ -15,20 +24,23 @@ def work_iterator(it):
1524

1625

1726
class ItertoolsThreading(unittest.TestCase):
18-
1927
@threading_helper.reap_threads
2028
def test_accumulate(self):
2129
number_of_iterations = 10
2230
for _ in range(number_of_iterations):
2331
it = accumulate(tuple(range(40)))
24-
threading_helper.run_concurrently(work_iterator, nthreads=10, args=[it])
32+
threading_helper.run_concurrently(
33+
work_iterator, nthreads=10, args=[it]
34+
)
2535

2636
@threading_helper.reap_threads
2737
def test_batched(self):
2838
number_of_iterations = 10
2939
for _ in range(number_of_iterations):
3040
it = batched(tuple(range(1000)), 2)
31-
threading_helper.run_concurrently(work_iterator, nthreads=10, args=[it])
41+
threading_helper.run_concurrently(
42+
work_iterator, nthreads=10, args=[it]
43+
)
3244

3345
@threading_helper.reap_threads
3446
def test_cycle(self):
@@ -46,28 +58,88 @@ def test_chain(self):
4658
number_of_iterations = 10
4759
for _ in range(number_of_iterations):
4860
it = chain(*[(1,)] * 200)
49-
threading_helper.run_concurrently(work_iterator, nthreads=6, args=[it])
61+
threading_helper.run_concurrently(
62+
work_iterator, nthreads=6, args=[it]
63+
)
5064

5165
@threading_helper.reap_threads
5266
def test_combinations_with_replacement(self):
5367
number_of_iterations = 6
5468
for _ in range(number_of_iterations):
5569
it = combinations_with_replacement(tuple(range(2)), 2)
56-
threading_helper.run_concurrently(work_iterator, nthreads=6, args=[it])
70+
threading_helper.run_concurrently(
71+
work_iterator, nthreads=6, args=[it]
72+
)
5773

5874
@threading_helper.reap_threads
5975
def test_permutations(self):
6076
number_of_iterations = 6
6177
for _ in range(number_of_iterations):
6278
it = permutations(tuple(range(4)), 2)
63-
threading_helper.run_concurrently(work_iterator, nthreads=6, args=[it])
79+
threading_helper.run_concurrently(
80+
work_iterator, nthreads=6, args=[it]
81+
)
6482

6583
@threading_helper.reap_threads
6684
def test_zip_longest(self):
6785
number_of_iterations = 10
6886
for _ in range(number_of_iterations):
6987
it = zip_longest(list(range(4)), list(range(8)), fillvalue=0)
70-
threading_helper.run_concurrently(work_iterator, nthreads=10, args=[it])
88+
threading_helper.run_concurrently(
89+
work_iterator, nthreads=10, args=[it]
90+
)
91+
92+
93+
class TestTeeConcurrent(unittest.TestCase):
94+
# itertools.tee branches share a linked list of internal data cells.
95+
# Concurrent iteration must not corrupt that shared state or crash the
96+
# free-threaded build. A crash shows up as the interpreter dying (not as a
97+
# caught exception); tee is documented as not thread-safe, so a
98+
# ``RuntimeError`` from the re-entrancy guard is an allowed outcome and is
99+
# tolerated here.
100+
101+
def test_same_branch(self):
102+
# Many threads consume the same tee branch.
103+
errors = []
104+
105+
def consume(it):
106+
try:
107+
for _ in it:
108+
pass
109+
except RuntimeError:
110+
pass
111+
except Exception as e:
112+
errors.append(e)
113+
114+
for _ in range(100):
115+
a, _ = tee(iter(range(2000)), 2)
116+
threading_helper.run_concurrently(consume, nthreads=8, args=(a,))
117+
118+
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
119+
120+
def test_sibling_branches(self):
121+
# Each thread consumes a different sibling branch of the same tee.
122+
errors = []
123+
124+
def make_worker(it):
125+
def consume():
126+
try:
127+
for _ in it:
128+
pass
129+
except RuntimeError:
130+
pass
131+
except Exception as e:
132+
errors.append(e)
133+
134+
return consume
135+
136+
for _ in range(100):
137+
branches = tee(iter(range(4000)), 8)
138+
threading_helper.run_concurrently(
139+
[make_worker(it) for it in branches]
140+
)
141+
142+
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
71143

72144

73145
if __name__ == "__main__":
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
Fix a crash when concurrently iterating an :func:`itertools.tee` iterator on
2+
the free-threaded build.

Modules/itertoolsmodule.c

Lines changed: 75 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -768,13 +768,17 @@ teedataobject_newinternal(itertools_state *state, PyObject *it)
768768
static PyObject *
769769
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
770770
{
771+
PyObject *link;
772+
Py_BEGIN_CRITICAL_SECTION(tdo);
771773
if (tdo->nextlink == NULL)
772774
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
773-
return Py_XNewRef(tdo->nextlink);
775+
link = Py_XNewRef(tdo->nextlink);
776+
Py_END_CRITICAL_SECTION();
777+
return link;
774778
}
775779

776780
static PyObject *
777-
teedataobject_getitem(teedataobject *tdo, int i)
781+
teedataobject_getitem_lock_held(teedataobject *tdo, int i)
778782
{
779783
PyObject *value;
780784

@@ -800,6 +804,16 @@ teedataobject_getitem(teedataobject *tdo, int i)
800804
return Py_NewRef(value);
801805
}
802806

807+
static PyObject *
808+
teedataobject_getitem(teedataobject *tdo, int i)
809+
{
810+
PyObject *result;
811+
Py_BEGIN_CRITICAL_SECTION(tdo);
812+
result = teedataobject_getitem_lock_held(tdo, i);
813+
Py_END_CRITICAL_SECTION();
814+
return result;
815+
}
816+
803817
static int
804818
teedataobject_traverse(PyObject *op, visitproc visit, void * arg)
805819
{
@@ -819,8 +833,11 @@ teedataobject_safe_decref(PyObject *obj)
819833
{
820834
while (obj && _PyObject_IsUniquelyReferenced(obj)) {
821835
teedataobject *tmp = teedataobject_CAST(obj);
822-
PyObject *nextlink = tmp->nextlink;
836+
PyObject *nextlink;
837+
Py_BEGIN_CRITICAL_SECTION(obj);
838+
nextlink = tmp->nextlink;
823839
tmp->nextlink = NULL;
840+
Py_END_CRITICAL_SECTION();
824841
Py_SETREF(obj, nextlink);
825842
}
826843
Py_XDECREF(obj);
@@ -833,11 +850,13 @@ teedataobject_clear(PyObject *op)
833850
PyObject *tmp;
834851
teedataobject *tdo = teedataobject_CAST(op);
835852

853+
Py_BEGIN_CRITICAL_SECTION(op);
836854
Py_CLEAR(tdo->it);
837855
for (i=0 ; i<tdo->numread ; i++)
838856
Py_CLEAR(tdo->values[i]);
839857
tmp = tdo->nextlink;
840858
tdo->nextlink = NULL;
859+
Py_END_CRITICAL_SECTION();
841860
teedataobject_safe_decref(tmp);
842861
return 0;
843862
}
@@ -930,20 +949,67 @@ static PyObject *
930949
tee_next(PyObject *op)
931950
{
932951
teeobject *to = teeobject_CAST(op);
933-
PyObject *value, *link;
952+
PyObject *value;
934953

954+
#ifndef Py_GIL_DISABLED
955+
/* The GIL already serializes access, so keep the simple path without the
956+
snapshot and revalidation that the free-threaded build needs. */
935957
if (to->index >= LINKCELLS) {
936-
link = teedataobject_jumplink(to->state, to->dataobj);
937-
if (link == NULL)
958+
PyObject *link = teedataobject_jumplink(to->state, to->dataobj);
959+
if (link == NULL) {
938960
return NULL;
961+
}
939962
Py_SETREF(to->dataobj, (teedataobject *)link);
940963
to->index = 0;
941964
}
942965
value = teedataobject_getitem(to->dataobj, to->index);
943-
if (value == NULL)
966+
if (value == NULL) {
944967
return NULL;
968+
}
945969
to->index++;
946970
return value;
971+
#else
972+
for (;;) {
973+
teedataobject *dataobj;
974+
int index;
975+
976+
/* Snapshot the branch position (strong ref to the shared data object)
977+
under the tee lock; the data object is locked separately, not nested,
978+
then the advance is revalidated. */
979+
Py_BEGIN_CRITICAL_SECTION(op);
980+
dataobj = (teedataobject *)Py_NewRef((PyObject *)to->dataobj);
981+
index = to->index;
982+
Py_END_CRITICAL_SECTION();
983+
984+
if (index < LINKCELLS) {
985+
value = teedataobject_getitem(dataobj, index);
986+
if (value != NULL) {
987+
Py_BEGIN_CRITICAL_SECTION(op);
988+
if (to->dataobj == dataobj && to->index == index) {
989+
to->index = index + 1;
990+
}
991+
Py_END_CRITICAL_SECTION();
992+
}
993+
Py_DECREF(dataobj);
994+
return value;
995+
}
996+
997+
PyObject *link = teedataobject_jumplink(to->state, dataobj);
998+
if (link == NULL) {
999+
Py_DECREF(dataobj);
1000+
return NULL;
1001+
}
1002+
Py_BEGIN_CRITICAL_SECTION(op);
1003+
if (to->dataobj == dataobj) {
1004+
Py_SETREF(to->dataobj, (teedataobject *)link);
1005+
to->index = 0;
1006+
link = NULL;
1007+
}
1008+
Py_END_CRITICAL_SECTION();
1009+
Py_XDECREF(link);
1010+
Py_DECREF(dataobj);
1011+
}
1012+
#endif
9471013
}
9481014

9491015
static int
@@ -962,8 +1028,10 @@ tee_copy_impl(teeobject *to)
9621028
if (newto == NULL) {
9631029
return NULL;
9641030
}
1031+
Py_BEGIN_CRITICAL_SECTION(to);
9651032
newto->dataobj = (teedataobject *)Py_NewRef(to->dataobj);
9661033
newto->index = to->index;
1034+
Py_END_CRITICAL_SECTION();
9671035
newto->weakreflist = NULL;
9681036
newto->state = to->state;
9691037
PyObject_GC_Track(newto);

0 commit comments

Comments
 (0)