import unittest
-from itertools import accumulate, batched, chain, combinations_with_replacement, cycle, permutations, zip_longest
+from itertools import (
+ accumulate,
+ batched,
+ chain,
+ combinations_with_replacement,
+ cycle,
+ permutations,
+ tee,
+ zip_longest,
+)
from test.support import threading_helper
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.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.run_concurrently(
+ work_iterator, nthreads=10, args=[it]
+ )
@threading_helper.reap_threads
def test_cycle(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.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.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.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])
+ 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__":
static PyObject *
teedataobject_jumplink(itertools_state *state, teedataobject *tdo)
{
+ PyObject *link;
+ Py_BEGIN_CRITICAL_SECTION(tdo);
if (tdo->nextlink == NULL)
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
- return Py_XNewRef(tdo->nextlink);
+ link = Py_XNewRef(tdo->nextlink);
+ Py_END_CRITICAL_SECTION();
+ return link;
}
static PyObject *
-teedataobject_getitem(teedataobject *tdo, int i)
+teedataobject_getitem_lock_held(teedataobject *tdo, int i)
{
PyObject *value;
return Py_NewRef(value);
}
+static PyObject *
+teedataobject_getitem(teedataobject *tdo, int i)
+{
+ PyObject *result;
+ Py_BEGIN_CRITICAL_SECTION(tdo);
+ result = teedataobject_getitem_lock_held(tdo, i);
+ Py_END_CRITICAL_SECTION();
+ return result;
+}
+
static int
teedataobject_traverse(PyObject *op, visitproc visit, void * arg)
{
{
while (obj && _PyObject_IsUniquelyReferenced(obj)) {
teedataobject *tmp = teedataobject_CAST(obj);
- PyObject *nextlink = tmp->nextlink;
+ PyObject *nextlink;
+ Py_BEGIN_CRITICAL_SECTION(obj);
+ nextlink = tmp->nextlink;
tmp->nextlink = NULL;
+ Py_END_CRITICAL_SECTION();
Py_SETREF(obj, nextlink);
}
Py_XDECREF(obj);
PyObject *tmp;
teedataobject *tdo = teedataobject_CAST(op);
+ Py_BEGIN_CRITICAL_SECTION(op);
Py_CLEAR(tdo->it);
for (i=0 ; i<tdo->numread ; i++)
Py_CLEAR(tdo->values[i]);
tmp = tdo->nextlink;
tdo->nextlink = NULL;
+ Py_END_CRITICAL_SECTION();
teedataobject_safe_decref(tmp);
return 0;
}
tee_next(PyObject *op)
{
teeobject *to = teeobject_CAST(op);
- PyObject *value, *link;
+ PyObject *value;
+#ifndef Py_GIL_DISABLED
+ /* The GIL already serializes access, so keep the simple path without the
+ snapshot and revalidation that the free-threaded build needs. */
if (to->index >= LINKCELLS) {
- link = teedataobject_jumplink(to->state, to->dataobj);
- if (link == NULL)
+ PyObject *link = teedataobject_jumplink(to->state, to->dataobj);
+ if (link == NULL) {
return NULL;
+ }
Py_SETREF(to->dataobj, (teedataobject *)link);
to->index = 0;
}
value = teedataobject_getitem(to->dataobj, to->index);
- if (value == NULL)
+ if (value == NULL) {
return NULL;
+ }
to->index++;
return value;
+#else
+ for (;;) {
+ teedataobject *dataobj;
+ int index;
+
+ /* Snapshot the branch position (strong ref to the shared data object)
+ under the tee lock; the data object is locked separately, not nested,
+ then the advance is revalidated. */
+ Py_BEGIN_CRITICAL_SECTION(op);
+ dataobj = (teedataobject *)Py_NewRef((PyObject *)to->dataobj);
+ index = to->index;
+ Py_END_CRITICAL_SECTION();
+
+ if (index < LINKCELLS) {
+ value = teedataobject_getitem(dataobj, index);
+ if (value != NULL) {
+ Py_BEGIN_CRITICAL_SECTION(op);
+ if (to->dataobj == dataobj && to->index == index) {
+ to->index = index + 1;
+ }
+ Py_END_CRITICAL_SECTION();
+ }
+ Py_DECREF(dataobj);
+ return value;
+ }
+
+ PyObject *link = teedataobject_jumplink(to->state, dataobj);
+ if (link == NULL) {
+ Py_DECREF(dataobj);
+ return NULL;
+ }
+ Py_BEGIN_CRITICAL_SECTION(op);
+ if (to->dataobj == dataobj) {
+ Py_SETREF(to->dataobj, (teedataobject *)link);
+ to->index = 0;
+ link = NULL;
+ }
+ Py_END_CRITICAL_SECTION();
+ Py_XDECREF(link);
+ Py_DECREF(dataobj);
+ }
+#endif
}
static int
if (newto == NULL) {
return NULL;
}
+ Py_BEGIN_CRITICAL_SECTION(to);
newto->dataobj = (teedataobject *)Py_NewRef(to->dataobj);
newto->index = to->index;
+ Py_END_CRITICAL_SECTION();
newto->weakreflist = NULL;
newto->state = to->state;
PyObject_GC_Track(newto);
projects / thirdparty / Python / cpython.git / commitdiff
