bpo-34410: Fix a crash in the tee iterator when re-enter it. (GH-15625)

RuntimeError is now raised in this case.
(cherry picked from commit 526a01467b3277f9fcf7f91e66c23321caa1245d)

Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 959424ff914390e8151089da5dfaa24a8aa9433e..76899ce753685a1ce8b37a9b11a720a184fa719a 100644 (file)
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -628,6 +628,10 @@ loops that truncate the stream.
    used anywhere else; otherwise, the *iterable* could get advanced without
    the tee objects being informed.
 
+   ``tee`` iterators are not threadsafe. A :exc:`RuntimeError` may be
+   raised when using simultaneously iterators returned by the same :func:`tee`
+   call, even if the original *iterable* is threadsafe.
+
    This itertool may require significant auxiliary storage (depending on how
    much temporary data needs to be stored). In general, if one iterator uses
    most or all of the data before another iterator starts, it is faster to use

diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py
index cbbb4c4f71d3b8d2e8fad564078ee14b0c467291..721a17556f7ed7da349e1034c08dd17c2ad6ff99 100644 (file)
--- a/Lib/test/test_itertools.py
+++ b/Lib/test/test_itertools.py
@@ -11,6 +11,7 @@ import pickle
 from functools import reduce
 import sys
 import struct
+import threading
 maxsize = support.MAX_Py_ssize_t
 minsize = -maxsize-1
 
@@ -1476,6 +1477,42 @@ class TestBasicOps(unittest.TestCase):
             del forward, backward
             raise
 
+    def test_tee_reenter(self):
+        class I:
+            first = True
+            def __iter__(self):
+                return self
+            def __next__(self):
+                first = self.first
+                self.first = False
+                if first:
+                    return next(b)
+
+        a, b = tee(I())
+        with self.assertRaisesRegex(RuntimeError, "tee"):
+            next(a)
+
+    def test_tee_concurrent(self):
+        start = threading.Event()
+        finish = threading.Event()
+        class I:
+            def __iter__(self):
+                return self
+            def __next__(self):
+                start.set()
+                finish.wait()
+
+        a, b = tee(I())
+        thread = threading.Thread(target=next, args=[a])
+        thread.start()
+        try:
+            start.wait()
+            with self.assertRaisesRegex(RuntimeError, "tee"):
+                next(b)
+        finally:
+            finish.set()
+            thread.join()
+
     def test_StopIteration(self):
         self.assertRaises(StopIteration, next, zip())
 

diff --git a/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst
new file mode 100644 (file)
index 0000000..64e778e
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2019-08-31-01-52-59.bpo-34410.7KbWZQ.rst
@@ -0,0 +1,2 @@
+Fixed a crash in the :func:`tee` iterator when re-enter it. RuntimeError is
+now raised in this case.

diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c
index bff9df608304b67a69acb4c7d61d242f63be0a17..0ddd98dea124748c5bbdc6f6c6f632b59df16229 100644 (file)
--- a/Modules/itertoolsmodule.c
+++ b/Modules/itertoolsmodule.c
@@ -395,6 +395,7 @@ typedef struct {
     PyObject_HEAD
     PyObject *it;
     int numread;                /* 0 <= numread <= LINKCELLS */
+    int running;
     PyObject *nextlink;
     PyObject *(values[LINKCELLS]);
 } teedataobject;
@@ -417,6 +418,7 @@ teedataobject_newinternal(PyObject *it)
     if (tdo == NULL)
         return NULL;
 
+    tdo->running = 0;
     tdo->numread = 0;
     tdo->nextlink = NULL;
     Py_INCREF(it);
@@ -445,7 +447,14 @@ teedataobject_getitem(teedataobject *tdo, int i)
     else {
         /* this is the lead iterator, so fetch more data */
         assert(i == tdo->numread);
+        if (tdo->running) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "cannot re-enter the tee iterator");
+            return NULL;
+        }
+        tdo->running = 1;
         value = PyIter_Next(tdo->it);
+        tdo->running = 0;
         if (value == NULL)
             return NULL;
         tdo->numread++;