Skip to content

Commit 8ea39d9

Browse files
tonghuarootnascheme
andcommitted
[3.14] 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. (cherry picked from commit 6c389c4) Co-authored-by: tonghuaroot (童话) <tonghuaroot@gmail.com> Co-authored-by: Neil Schemenauer <nas-github@arctrix.com>
1 parent 6374772 commit 8ea39d9

3 files changed

Lines changed: 141 additions & 7 deletions

File tree

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import unittest
2+
from itertools import (
3+
tee,
4+
)
5+
from test.support import threading_helper
6+
7+
8+
threading_helper.requires_working_threading(module=True)
9+
10+
11+
class TestTeeConcurrent(unittest.TestCase):
12+
# itertools.tee branches share a linked list of internal data cells.
13+
# Concurrent iteration must not corrupt that shared state or crash the
14+
# free-threaded build. A crash shows up as the interpreter dying (not as a
15+
# caught exception); tee is documented as not thread-safe, so a
16+
# ``RuntimeError`` from the re-entrancy guard is an allowed outcome and is
17+
# tolerated here.
18+
19+
def test_same_branch(self):
20+
# Many threads consume the same tee branch.
21+
errors = []
22+
23+
def consume(it):
24+
try:
25+
for _ in it:
26+
pass
27+
except RuntimeError:
28+
pass
29+
except Exception as e:
30+
errors.append(e)
31+
32+
for _ in range(100):
33+
a, _ = tee(iter(range(2000)), 2)
34+
threading_helper.run_concurrently(consume, nthreads=8, args=(a,))
35+
36+
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
37+
38+
def test_sibling_branches(self):
39+
# Each thread consumes a different sibling branch of the same tee.
40+
errors = []
41+
42+
def make_worker(it):
43+
def consume():
44+
try:
45+
for _ in it:
46+
pass
47+
except RuntimeError:
48+
pass
49+
except Exception as e:
50+
errors.append(e)
51+
52+
return consume
53+
54+
for _ in range(100):
55+
branches = tee(iter(range(4000)), 8)
56+
threading_helper.run_concurrently(
57+
[make_worker(it) for it in branches]
58+
)
59+
60+
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
61+
62+
63+
if __name__ == "__main__":
64+
unittest.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
@@ -772,13 +772,17 @@ teedataobject_newinternal(itertools_state *state, PyObject *it)
772772
static PyObject *
773773
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
774774
{
775+
PyObject *link;
776+
Py_BEGIN_CRITICAL_SECTION(tdo);
775777
if (tdo->nextlink == NULL)
776778
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
777-
return Py_XNewRef(tdo->nextlink);
779+
link = Py_XNewRef(tdo->nextlink);
780+
Py_END_CRITICAL_SECTION();
781+
return link;
778782
}
779783

780784
static PyObject *
781-
teedataobject_getitem(teedataobject *tdo, int i)
785+
teedataobject_getitem_lock_held(teedataobject *tdo, int i)
782786
{
783787
PyObject *value;
784788

@@ -804,6 +808,16 @@ teedataobject_getitem(teedataobject *tdo, int i)
804808
return Py_NewRef(value);
805809
}
806810

811+
static PyObject *
812+
teedataobject_getitem(teedataobject *tdo, int i)
813+
{
814+
PyObject *result;
815+
Py_BEGIN_CRITICAL_SECTION(tdo);
816+
result = teedataobject_getitem_lock_held(tdo, i);
817+
Py_END_CRITICAL_SECTION();
818+
return result;
819+
}
820+
807821
static int
808822
teedataobject_traverse(PyObject *op, visitproc visit, void * arg)
809823
{
@@ -823,8 +837,11 @@ teedataobject_safe_decref(PyObject *obj)
823837
{
824838
while (obj && _PyObject_IsUniquelyReferenced(obj)) {
825839
teedataobject *tmp = teedataobject_CAST(obj);
826-
PyObject *nextlink = tmp->nextlink;
840+
PyObject *nextlink;
841+
Py_BEGIN_CRITICAL_SECTION(obj);
842+
nextlink = tmp->nextlink;
827843
tmp->nextlink = NULL;
844+
Py_END_CRITICAL_SECTION();
828845
Py_SETREF(obj, nextlink);
829846
}
830847
Py_XDECREF(obj);
@@ -837,11 +854,13 @@ teedataobject_clear(PyObject *op)
837854
PyObject *tmp;
838855
teedataobject *tdo = teedataobject_CAST(op);
839856

857+
Py_BEGIN_CRITICAL_SECTION(op);
840858
Py_CLEAR(tdo->it);
841859
for (i=0 ; i<tdo->numread ; i++)
842860
Py_CLEAR(tdo->values[i]);
843861
tmp = tdo->nextlink;
844862
tdo->nextlink = NULL;
863+
Py_END_CRITICAL_SECTION();
845864
teedataobject_safe_decref(tmp);
846865
return 0;
847866
}
@@ -934,20 +953,67 @@ static PyObject *
934953
tee_next(PyObject *op)
935954
{
936955
teeobject *to = teeobject_CAST(op);
937-
PyObject *value, *link;
956+
PyObject *value;
938957

958+
#ifndef Py_GIL_DISABLED
959+
/* The GIL already serializes access, so keep the simple path without the
960+
snapshot and revalidation that the free-threaded build needs. */
939961
if (to->index >= LINKCELLS) {
940-
link = teedataobject_jumplink(to->state, to->dataobj);
941-
if (link == NULL)
962+
PyObject *link = teedataobject_jumplink(to->state, to->dataobj);
963+
if (link == NULL) {
942964
return NULL;
965+
}
943966
Py_SETREF(to->dataobj, (teedataobject *)link);
944967
to->index = 0;
945968
}
946969
value = teedataobject_getitem(to->dataobj, to->index);
947-
if (value == NULL)
970+
if (value == NULL) {
948971
return NULL;
972+
}
949973
to->index++;
950974
return value;
975+
#else
976+
for (;;) {
977+
teedataobject *dataobj;
978+
int index;
979+
980+
/* Snapshot the branch position (strong ref to the shared data object)
981+
under the tee lock; the data object is locked separately, not nested,
982+
then the advance is revalidated. */
983+
Py_BEGIN_CRITICAL_SECTION(op);
984+
dataobj = (teedataobject *)Py_NewRef((PyObject *)to->dataobj);
985+
index = to->index;
986+
Py_END_CRITICAL_SECTION();
987+
988+
if (index < LINKCELLS) {
989+
value = teedataobject_getitem(dataobj, index);
990+
if (value != NULL) {
991+
Py_BEGIN_CRITICAL_SECTION(op);
992+
if (to->dataobj == dataobj && to->index == index) {
993+
to->index = index + 1;
994+
}
995+
Py_END_CRITICAL_SECTION();
996+
}
997+
Py_DECREF(dataobj);
998+
return value;
999+
}
1000+
1001+
PyObject *link = teedataobject_jumplink(to->state, dataobj);
1002+
if (link == NULL) {
1003+
Py_DECREF(dataobj);
1004+
return NULL;
1005+
}
1006+
Py_BEGIN_CRITICAL_SECTION(op);
1007+
if (to->dataobj == dataobj) {
1008+
Py_SETREF(to->dataobj, (teedataobject *)link);
1009+
to->index = 0;
1010+
link = NULL;
1011+
}
1012+
Py_END_CRITICAL_SECTION();
1013+
Py_XDECREF(link);
1014+
Py_DECREF(dataobj);
1015+
}
1016+
#endif
9511017
}
9521018

9531019
static int
@@ -966,8 +1032,10 @@ tee_copy_impl(teeobject *to)
9661032
if (newto == NULL) {
9671033
return NULL;
9681034
}
1035+
Py_BEGIN_CRITICAL_SECTION(to);
9691036
newto->dataobj = (teedataobject *)Py_NewRef(to->dataobj);
9701037
newto->index = to->index;
1038+
Py_END_CRITICAL_SECTION();
9711039
newto->weakreflist = NULL;
9721040
newto->state = to->state;
9731041
PyObject_GC_Track(newto);

0 commit comments

Comments
 (0)