mirror of
https://github.com/python/cpython.git
synced 2026-07-31 14:16:33 -04:00
6c389c446e
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 <[email protected]>
147 lines
4.3 KiB
Python
147 lines
4.3 KiB
Python
import unittest
|
|
from itertools import (
|
|
accumulate,
|
|
batched,
|
|
chain,
|
|
combinations_with_replacement,
|
|
cycle,
|
|
permutations,
|
|
tee,
|
|
zip_longest,
|
|
)
|
|
from test.support import threading_helper
|
|
|
|
|
|
threading_helper.requires_working_threading(module=True)
|
|
|
|
|
|
def work_iterator(it):
|
|
while True:
|
|
try:
|
|
next(it)
|
|
except StopIteration:
|
|
break
|
|
|
|
|
|
class ItertoolsThreading(unittest.TestCase):
|
|
@threading_helper.reap_threads
|
|
def test_accumulate(self):
|
|
number_of_iterations = 10
|
|
for _ in range(number_of_iterations):
|
|
it = accumulate(tuple(range(40)))
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=10, args=[it]
|
|
)
|
|
|
|
@threading_helper.reap_threads
|
|
def test_batched(self):
|
|
number_of_iterations = 10
|
|
for _ in range(number_of_iterations):
|
|
it = batched(tuple(range(1000)), 2)
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=10, args=[it]
|
|
)
|
|
|
|
@threading_helper.reap_threads
|
|
def test_cycle(self):
|
|
def work(it):
|
|
for _ in range(400):
|
|
next(it)
|
|
|
|
number_of_iterations = 6
|
|
for _ in range(number_of_iterations):
|
|
it = cycle((1, 2, 3, 4))
|
|
threading_helper.run_concurrently(work, nthreads=6, args=[it])
|
|
|
|
@threading_helper.reap_threads
|
|
def test_chain(self):
|
|
number_of_iterations = 10
|
|
for _ in range(number_of_iterations):
|
|
it = chain(*[(1,)] * 200)
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=6, args=[it]
|
|
)
|
|
|
|
@threading_helper.reap_threads
|
|
def test_combinations_with_replacement(self):
|
|
number_of_iterations = 6
|
|
for _ in range(number_of_iterations):
|
|
it = combinations_with_replacement(tuple(range(2)), 2)
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=6, args=[it]
|
|
)
|
|
|
|
@threading_helper.reap_threads
|
|
def test_permutations(self):
|
|
number_of_iterations = 6
|
|
for _ in range(number_of_iterations):
|
|
it = permutations(tuple(range(4)), 2)
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=6, args=[it]
|
|
)
|
|
|
|
@threading_helper.reap_threads
|
|
def test_zip_longest(self):
|
|
number_of_iterations = 10
|
|
for _ in range(number_of_iterations):
|
|
it = zip_longest(list(range(4)), list(range(8)), fillvalue=0)
|
|
threading_helper.run_concurrently(
|
|
work_iterator, nthreads=10, args=[it]
|
|
)
|
|
|
|
|
|
class TestTeeConcurrent(unittest.TestCase):
|
|
# itertools.tee branches share a linked list of internal data cells.
|
|
# Concurrent iteration must not corrupt that shared state or crash the
|
|
# free-threaded build. A crash shows up as the interpreter dying (not as a
|
|
# caught exception); tee is documented as not thread-safe, so a
|
|
# ``RuntimeError`` from the re-entrancy guard is an allowed outcome and is
|
|
# tolerated here.
|
|
|
|
def test_same_branch(self):
|
|
# Many threads consume the same tee branch.
|
|
errors = []
|
|
|
|
def consume(it):
|
|
try:
|
|
for _ in it:
|
|
pass
|
|
except RuntimeError:
|
|
pass
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
for _ in range(100):
|
|
a, _ = tee(iter(range(2000)), 2)
|
|
threading_helper.run_concurrently(consume, nthreads=8, args=(a,))
|
|
|
|
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
|
|
|
|
def test_sibling_branches(self):
|
|
# Each thread consumes a different sibling branch of the same tee.
|
|
errors = []
|
|
|
|
def make_worker(it):
|
|
def consume():
|
|
try:
|
|
for _ in it:
|
|
pass
|
|
except RuntimeError:
|
|
pass
|
|
except Exception as e:
|
|
errors.append(e)
|
|
|
|
return consume
|
|
|
|
for _ in range(100):
|
|
branches = tee(iter(range(4000)), 8)
|
|
threading_helper.run_concurrently(
|
|
[make_worker(it) for it in branches]
|
|
)
|
|
|
|
self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|