itertools.tee() 的结果是线程安全的（Python）

Question

假设我有这样的Python代码：

from itertools import count, tee
original = count()     # just an example, can be another iterable
a, b = tee(original)

问题是，如果我开始在一个线程中迭代“a”，同时在另一个线程中迭代“b”，会出现任何问题吗？显然，a 和 b 共享一些数据（原始的可迭代数据，+一些额外的东西，内部缓冲区或其他东西）。那么，a.next() 和 b.next() 在访问此共享数据时会执行适当的锁定吗？

Answer 1

更新！ 由 tee 引起的段错误已在 python 2.7、3.7、3.8 及以上版本的最新版本中得到修复。为了线程安全，你仍然需要自己管理并发访问，你可以使用我下面的解决方案。

tl;dr

In CPython, `itertools.tee` is thread safe *if and only if* the original iterator is implemented in C/C++, i.e. doesn't use **any** python.

如果原始迭代器 it 是用 Python 编写的，例如类实例或生成器，则 itertools.tee(it) 是不是 线程安全。在最好的情况下，你只会得到一个异常（你会的），而在最坏的情况下，Python 会崩溃。

这里没有使用 tee，而是一个线程安全的包装类和函数：

class safeteeobject(object):
    """tee object wrapped to make it thread-safe"""
    def __init__(self, teeobj, lock):
        self.teeobj = teeobj
        self.lock = lock
    def __iter__(self):
        return self
    def __next__(self):
        with self.lock:
            return next(self.teeobj)
    def __copy__(self):
        return safeteeobject(self.teeobj.__copy__(), self.lock)

def safetee(iterable, n=2):
    """tuple of n independent thread-safe iterators"""
    lock = Lock()
    return tuple(safeteeobject(teeobj, lock) for teeobj in tee(iterable, n))

我现在将（大量）扩展 tee 何时是线程以及何时不是线程-安全，以及为什么。

可以的示例

Let's run some code (this is python 3 code, for python 2 use `itertools.izip` instead of `zip` to have same behavior):

>>> from itertools import tee, count
>>> from threading import Thread

>>> def limited_sum(it):
...     s = 0
...     for elem, _ in zip(it, range(1000000)):
...         s += elem
...     print(elem)

>>> a, b = tee(count())
>>> [Thread(target=limited_sum, args=(it,)).start() for it in [a, b]]
# prints 499999500000 twice, which is in fact the same 1+...+999999

itertools.count 完全用 C 语言编写在 CPython 项目的 Modules/itertoolsmodule.c 文件中，因此它工作得很好。

对于列表、元组、集合、范围、字典（键、值和项）、collections.defaultdict（键、值和项）以及其他一些也是如此。

它不起作用的示例 - 生成器

A very short example is using a generator:

>>> gen = (i for i in range(1000000))
>>> a, b = tee(gen)
>>> [Thread(target=sum, args=(it,)).start() for it in [a, b]]

Exception in thread Thread-10:
Traceback (most recent call last):
  File "/usr/lib/python3.4/threading.py", line 920, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.4/threading.py", line 868, in run
    self._target(*self._args, **self._kwargs)
ValueError: generator already executing

是的，tee 是用 C 编写的，并且 GIL 确实一次执行一个字节代码。但上面的例子表明，这不足以保证线程安全。沿着这条线的某个地方发生了这样的事情：

1. 两个线程在其 tee_object 实例上调用 next 次数相同，
2. 线程 1 调用 next(a)，
3. 它需要获取一个新元素，因此线程 1 现在调用 next(gen)，
4. gen 是用 python 编写的。比如说，gen.__next__的第一个字节码CPython决定切换线程，
5. 线程2恢复并调用next(b)，
6. 它需要获取一个新元素，所以它调用 next(gen)
7. 由于 gen.__next__ 已经在线程 1 中运行，我们得到一个异常。

不起作用的示例 - 迭代器对象

Ok, maybe it's just not thread-safe to use generators inside `tee`. Then we run a variant of the above code that uses an iterator object:

>>> from itertools import tee
>>> from threading import Thread
>>> class countdown(object):
...     def __init__(self, n):
...         self.i = n
...     def __iter__(self):
...         return self
...     def __next__(self):
...         self.i -= 1
...         if self.i < 0:
...             raise StopIteration
...         return self.i
... 
>>> a, b = tee(countdown(100000))
>>> [Thread(target=sum, args=(it,)).start() for it in [a, b]]
Segmentation fault (core dumped)

上面的代码在 Ubuntu、Windows 7 和 OSX 上的 python 2.7.13 和 3.6（可能还有所有 cpython 版本）中崩溃。我现在还不想透露原因，先一步。

如果我在迭代器中使用锁会怎样？

Maybe the above code crashes because our iterator itself was not thread-safe. Let's add a lock and see what happens:

>>> from itertools import tee
>>> from threading import Thread, Lock
>>> class countdown(object):
...     def __init__(self, n):
...         self.i = n
...         self.lock = Lock()
...     def __iter__(self):
...         return self
...     def __next__(self):
...         with self.lock:
...             self.i -= 1
...             if self.i < 0:
...                 raise StopIteration
...             return self.i
... 
>>> a, b = tee(countdown(100000))
>>> [Thread(target=sum, args=(it,)).start() for it in [a, b]]
Segmentation fault (core dumped)

在迭代器中添加锁不足以使 tee 线程安全。

为什么tee不是线程安全的

问题的症结在于CPython的Modules/itertoolsmodule.c文件中teedataobject的getitem方法。 tee 的实现非常酷，通过优化可以节省 RAM 调用：tee 返回“tee 对象”，每个对象都保存对头 teedataobject 的引用。这些反过来就像链表中的链接，但不是保存单个元素 - 它们保存 57。这对于我们的目的来说并不重要，但它就是这样。以下是 teedataobject 的 getitem 函数：

static PyObject *
teedataobject_getitem(teedataobject *tdo, int i)
{
    PyObject *value;

    assert(i < LINKCELLS);
    if (i < tdo->numread)
        value = tdo->values[i];
    else {
        /* this is the lead iterator, so fetch more data */
        assert(i == tdo->numread);
        value = PyIter_Next(tdo->it);
        if (value == NULL)
            return NULL;
        tdo->numread++;
        tdo->values[i] = value;
    }
    Py_INCREF(value);
    return value;
}

当请求一个元素时，teedataobject 检查是否已准备好一个元素。如果是，则返回它。如果没有，那么它会在原始迭代器上调用 next 。如果迭代器是用 python 编写的，那么代码可能会挂起。所以问题是这样的：

1. 两个线程调用 next 的次数相同，
2. 线程 1 调用 next(a)，并且 C 代码到达 PyIter_Next< /code> 上面调用。比如说，在 next(gen) 的第一个字节码上，CPython 决定切换线程。
3. 线程 2 调用 next(b)，由于它仍然需要一个新元素，因此 C 代码进入 PyIter_Next 调用，

此时两个线程处于同一位置，i 和 tdo->numread 具有相同的值。请注意，tdo->numread 只是一个变量，用于跟踪 57 个单元格链接中 teedataobject 应写入下一个的位置。

4. 线程 2 完成对 PyIter_Next 的调用并返回一个元素。在某个时刻，CPython 决定再次切换线程，

5. 线程 1 恢复，完成对 PyIter_Next 的调用，然后运行以下两行：

    tdo->numread++;
        tdo->values[i] = 值；
   

6. 但是线程 2 已经设置了 tdo->values[i]!< /p>

这已经足以表明 tee 不是线程安全的，因为我们丢失了线程 2 放入 tdo->values[i] 中的值。但这并不能解释崩溃的原因。

假设 i 为 56。由于两个线程都调用 tdo->numread++，它现在达到 58 - 高于 57，即 tdo->values 的分配大小。线程 1 也继续移动后，对象 tdo 不再有引用，可以被删除了。这是 teedataobject 的清除函数：

static int
teedataobject_clear(teedataobject *tdo)
{
    int i;
    PyObject *tmp;

    Py_CLEAR(tdo->it);
    for (i=0 ; i<tdo->numread ; i++)
        Py_CLEAR(tdo->values[i]); // <----- PROBLEM!!!
    tmp = tdo->nextlink;
    tdo->nextlink = NULL;
    teedataobject_safe_decref(tmp);
    return 0;
}

在标记为“PROBLEM”的行，CPython 将尝试清除 tdo->values[57]。这就是崩溃发生的地方。嗯，有些时候。崩溃的地方不止一处，我只是想展示一个。

现在您知道了 - itertools.tee 不是线程安全的。

一种解决方案 - 外部锁

我们可以在 tee.__next__ 周围放置一个锁，而不是在迭代器的 __next__ 内部锁定。这意味着整个 teedataobject.__getitem__ 方法每次都会由单个线程调用。我在这个答案的开头给出了一个简短的实现。它是线程安全的 tee 的直接替代品。它唯一没有实现 tee 所做的事情 - 是酸洗。由于锁不可被picklable，所以添加它并不简单。但是，当然，这是可以做到的。

Answer 2

如果文档中显示等效代码，请参见：

• http://docs.python。 org/library/itertools.html#itertools.tee

是正确的，那么不，它不会是线程安全的。

请注意，尽管 deque 被记录为具有线程安全的追加和弹出，它不对使用它的代码做出任何保证。

由于主代码最终可能会向底层迭代器询问多个线程上的元素，因此您需要有一个线程安全的集合和迭代器作为输入，以便 tee 安全。

Answer 3

您可以使用此解决方法使 tee 线程安全：

import itertools
import threading
from typing import Iterable, Iterator, TypeVar

T = TypeVar('T')


class IterLock:
    def __init__(self):
        self.lock = threading.Lock()

    def __call__(self, it: Iterator[T]) -> Iterator[T]:
        def locked_iter():
            while True:
                with self.lock:
                    el = next(it)
                yield el
        return locked_iter()


def tee_safe(it: Iterable[T], n: int):
    lock = IterLock()
    return tuple(lock(sub_it) for sub_it in itertools.tee(it, n))

Answer 4

在 C-Python 中，itertools.tee() 及其返回的迭代器是用 C 代码实现的。这意味着 GIL 应该保护它不被多个线程同时调用。它可能会正常工作，并且不会使解释器崩溃，但不能保证线程安全。

简而言之，不要冒险。

Answer 5

我想分享一下我在 Python 3.6.9 和 3.7.4 环境中使用 itertools.tee 将大尺寸 plat 文本文件拆分为来自/到 s3 的多个 csv 文件的经验。

我的数据流来自s3 zipfile，s3fs读取iter，用于数据类转换的map iter，tee iter，用于数据类过滤器的map iter，循环iter并捕获数据并使用s3fs write和/或本地写入以csv格式写入s3和s3fs 放入 s3。

itertools.tee 在 zipfile 进程堆栈上失败。

上面，Dror Speiser，safetee 工作正常，但是由于数据集分布不好或处理延迟，tee 对象之间的任何不平衡都会导致内存使用量增加。
此外，它无法正常使用多处理日志记录，可能与此错误有关：https://bugs.python。 org/issue34410

下面的代码只是在 tee 对象之间添加简单的流程控制，以防止内存增加和 OOM Killer 情况。

希望对以后的参考有所帮助。

import time
import threading
import logging
from itertools import tee
from collections import Counter

logger = logging.getLogger(__name__)


FLOW_WAIT_GAP = 1000  # flow gap for waiting
FLOW_WAIT_TIMEOUT = 60.0  # flow wait timeout


class Safetee:
    """tee object wrapped to make it thread-safe and flow controlled"""

    def __init__(self, teeobj, lock, flows, teeidx):
        self.teeobj = teeobj
        self.lock = lock
        self.flows = flows
        self.mykey = teeidx
        self.logcnt = 0

    def __iter__(self):
        return self

    def __next__(self):
        waitsec = 0.0
        while True:
            with self.lock:
                flowgap = self.flows[self.mykey] - self.flows[len(self.flows) - 1]
                if flowgap < FLOW_WAIT_GAP or waitsec > FLOW_WAIT_TIMEOUT:
                    nextdata = next(self.teeobj)
                    self.flows[self.mykey] += 1
                    return nextdata

            waitthis = min(flowgap / FLOW_WAIT_GAP, FLOW_WAIT_TIMEOUT / 3)
            waitsec += waitthis

            time.sleep(waitthis)

            if waitsec > FLOW_WAIT_TIMEOUT and self.logcnt < 5:
                self.logcnt += 1
                logger.debug(f'tee wait seconds={waitsec:.2f}, mykey={self.mykey}, flows={self.flows}')

    def __copy__(self):
        return Safetee(self.teeobj.__copy__(), self.lock, self.flows, self.teeidx)


def safetee(iterable, n=2):
    """tuple of n independent thread-safe and flow controlled iterators"""
    lock = threading.Lock()
    flows = Counter()
    return tuple(Safetee(teeobj, lock, flows, teeidx) for teeidx, teeobj in enumerate(tee(iterable, n)))