Skip to content

Commit 103c022

Browse files
eendebakptclaude
andcommitted
gh-149557, gh-123471: Make itertools.pairwise safe under concurrency and re-entry
Wrap pairwise.__next__ in a critical section so concurrent callers from different threads serialize cleanly in the free-threaded build, and add a re-entrancy guard so a same-thread re-entrant call (e.g. when the input iterator's __next__ calls back into pairwise.__next__) raises RuntimeError instead of crashing or producing garbled output. Matches the long-standing behavior of itertools.tee and of generators, and fixes the use-after-free crash reported in gh-149557. With re-entry blocked, po->old is safe as a borrowed reference for the duration of __next__, so the defensive incref/decref and re-read guard added by gh-109788 are no longer needed and have been removed. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent ef2246f commit 103c022

4 files changed

Lines changed: 82 additions & 73 deletions

File tree

Lib/test/test_free_threading/test_itertools.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import unittest
2-
from itertools import accumulate, batched, chain, combinations_with_replacement, cycle, permutations, zip_longest
2+
from itertools import accumulate, batched, chain, combinations_with_replacement, cycle, pairwise, permutations, zip_longest
33
from test.support import threading_helper
44

55

@@ -55,6 +55,13 @@ def test_combinations_with_replacement(self):
5555
it = combinations_with_replacement(tuple(range(2)), 2)
5656
threading_helper.run_concurrently(work_iterator, nthreads=6, args=[it])
5757

58+
@threading_helper.reap_threads
59+
def test_pairwise(self):
60+
number_of_iterations = 10
61+
for _ in range(number_of_iterations):
62+
it = pairwise(tuple(range(100)))
63+
threading_helper.run_concurrently(work_iterator, nthreads=10, args=[it])
64+
5865
@threading_helper.reap_threads
5966
def test_permutations(self):
6067
number_of_iterations = 6

Lib/test/test_itertools.py

Lines changed: 31 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -968,76 +968,49 @@ def test_pairwise(self):
968968
pairwise(None) # non-iterable argument
969969

970970
def test_pairwise_reenter(self):
971-
def check(reenter_at, expected):
971+
# gh-149557: re-entering pairwise.__next__ from within the input
972+
# iterator's __next__ now raises RuntimeError, matching the behavior
973+
# of itertools.tee and of generators.
974+
def check(reenter_at):
972975
class I:
973976
count = 0
974977
def __iter__(self):
975978
return self
976979
def __next__(self):
977-
self.count +=1
980+
self.count += 1
978981
if self.count in reenter_at:
979982
return next(it)
980983
return [self.count] # new object
981984

982985
it = pairwise(I())
983-
for item in expected:
984-
self.assertEqual(next(it), item)
985-
986-
check({1}, [
987-
(([2], [3]), [4]),
988-
([4], [5]),
989-
])
990-
check({2}, [
991-
([1], ([1], [3])),
992-
(([1], [3]), [4]),
993-
([4], [5]),
994-
])
995-
check({3}, [
996-
([1], [2]),
997-
([2], ([2], [4])),
998-
(([2], [4]), [5]),
999-
([5], [6]),
1000-
])
1001-
check({1, 2}, [
1002-
((([3], [4]), [5]), [6]),
1003-
([6], [7]),
1004-
])
1005-
check({1, 3}, [
1006-
(([2], ([2], [4])), [5]),
1007-
([5], [6]),
1008-
])
1009-
check({1, 4}, [
1010-
(([2], [3]), (([2], [3]), [5])),
1011-
((([2], [3]), [5]), [6]),
1012-
([6], [7]),
1013-
])
1014-
check({2, 3}, [
1015-
([1], ([1], ([1], [4]))),
1016-
(([1], ([1], [4])), [5]),
1017-
([5], [6]),
1018-
])
986+
with self.assertRaisesRegex(RuntimeError, "pairwise"):
987+
list(it)
1019988

1020-
def test_pairwise_reenter2(self):
1021-
def check(maxcount, expected):
1022-
class I:
1023-
count = 0
1024-
def __iter__(self):
1025-
return self
1026-
def __next__(self):
1027-
if self.count >= maxcount:
1028-
raise StopIteration
1029-
self.count +=1
1030-
if self.count == 1:
1031-
return next(it, None)
1032-
return [self.count] # new object
989+
# Re-entry on any of the first few __next__ calls must raise.
990+
for reenter_at in [{1}, {2}, {3}, {1, 2}, {1, 3}, {1, 4}, {2, 3}]:
991+
with self.subTest(reenter_at=reenter_at):
992+
check(reenter_at)
1033993

1034-
it = pairwise(I())
1035-
self.assertEqual(list(it), expected)
1036-
1037-
check(1, [])
1038-
check(2, [])
1039-
check(3, [])
1040-
check(4, [(([2], [3]), [4])])
994+
def test_pairwise_reenter2(self):
995+
# gh-149557: variant of test_pairwise_reenter where the re-entrant
996+
# call uses next(it, None) (i.e. supplies a default). Even with a
997+
# default that suppresses StopIteration, the re-entry itself must
998+
# raise RuntimeError, propagating out of the outer pairwise call.
999+
class I:
1000+
count = 0
1001+
def __iter__(self):
1002+
return self
1003+
def __next__(self):
1004+
if self.count >= 4:
1005+
raise StopIteration
1006+
self.count += 1
1007+
if self.count == 1:
1008+
return next(it, None)
1009+
return [self.count] # new object
1010+
1011+
it = pairwise(I())
1012+
with self.assertRaisesRegex(RuntimeError, "pairwise"):
1013+
list(it)
10411014

10421015
def test_product(self):
10431016
for args, result in [
Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
:class:`itertools.pairwise` iterators are now safe to share between
2+
threads in the :term:`free threaded <free threading>` build: concurrent
3+
calls to ``__next__`` are serialized via a critical section
4+
(:gh:`123471`). Re-entrant calls (e.g. when the input iterator's
5+
``__next__`` calls back into the :class:`!pairwise` object) now raise
6+
:exc:`RuntimeError`, matching the behavior of generators and
7+
:class:`itertools.tee` and fixing a use-after-free crash
8+
(:gh:`149557`). Patch by Pieter Eendebak.

Modules/itertoolsmodule.c

Lines changed: 35 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -278,6 +278,7 @@ typedef struct {
278278
PyObject *it;
279279
PyObject *old;
280280
PyObject *result;
281+
uint8_t running;
281282
} pairwiseobject;
282283

283284
#define pairwiseobject_CAST(op) ((pairwiseobject *)(op))
@@ -344,36 +345,44 @@ pairwise_traverse(PyObject *op, visitproc visit, void *arg)
344345
}
345346

346347
static PyObject *
347-
pairwise_next(PyObject *op)
348+
pairwise_next_lock_held(PyObject *op)
348349
{
349350
pairwiseobject *po = pairwiseobject_CAST(op);
351+
352+
// Re-entrancy guard. The enclosing critical section serializes calls
353+
// from different threads; this flag catches same-thread re-entry (the
354+
// input iterator's __next__ calling back into pairwise.__next__) and
355+
// raises, matching itertools.tee and generators.
356+
if (po->running) {
357+
PyErr_SetString(PyExc_RuntimeError,
358+
"pairwise() iterator already executing");
359+
return NULL;
360+
}
361+
po->running = 1;
362+
350363
PyObject *it = po->it;
351364
PyObject *old = po->old;
352365
PyObject *new, *result;
353366

354367
if (it == NULL) {
355-
return NULL;
368+
result = NULL;
369+
goto done;
356370
}
357371
if (old == NULL) {
358372
old = (*Py_TYPE(it)->tp_iternext)(it);
359-
Py_XSETREF(po->old, old);
360373
if (old == NULL) {
361374
Py_CLEAR(po->it);
362-
return NULL;
363-
}
364-
it = po->it;
365-
if (it == NULL) {
366-
Py_CLEAR(po->old);
367-
return NULL;
375+
result = NULL;
376+
goto done;
368377
}
378+
po->old = old; // po->old was NULL; no decref needed
369379
}
370-
Py_INCREF(old);
371380
new = (*Py_TYPE(it)->tp_iternext)(it);
372381
if (new == NULL) {
373382
Py_CLEAR(po->it);
374383
Py_CLEAR(po->old);
375-
Py_DECREF(old);
376-
return NULL;
384+
result = NULL;
385+
goto done;
377386
}
378387

379388
result = po->result;
@@ -393,8 +402,20 @@ pairwise_next(PyObject *op)
393402
result = _PyTuple_FromPair(old, new);
394403
}
395404

396-
Py_XSETREF(po->old, new);
397-
Py_DECREF(old);
405+
Py_SETREF(po->old, new); // po->old == old, known non-NULL
406+
407+
done:
408+
po->running = 0;
409+
return result;
410+
}
411+
412+
static PyObject *
413+
pairwise_next(PyObject *op)
414+
{
415+
PyObject *result;
416+
Py_BEGIN_CRITICAL_SECTION(op);
417+
result = pairwise_next_lock_held(op);
418+
Py_END_CRITICAL_SECTION();
398419
return result;
399420
}
400421

0 commit comments

Comments
 (0)