假设我有以下 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()是否会执行适当的锁定？
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是线程安全的，何时不是，以及原因。
>>> 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是在CPython项目的文件Modules/itertoolsmodule.c中完全使用C编写的，因此它可以正常工作。
同样适用于：lists、tuples、sets、range、dictionaries（键、值和项）、collections.defaultdict（键、值和项）和其他一些东西。
一个非常简短的示例是使用生成器：
>>> 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确实每次执行一个字节码。但上面的示例表明这还不足以确保线程安全。在某个地方发生了以下情况:
next,next(a),next(gen),gen是用Python编写的。在gen.__next__的第一个字节码中，CPython决定切换线程，next(b),next(gen)gen.__next__已经在线程1中运行，因此我们得到一个异常。>>> 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)
上述代码在Python 2.7.13和3.6（以及可能所有CPython版本）上，在Ubuntu、Windows 7和OSX上崩溃。我暂时不想透露原因，还有一步。
或许上述代码崩溃是因为迭代器本身不是线程安全的。我们添加一个锁看看会发生什么：
>>> 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线程安全。
问题的关键在于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。这就是问题所在：1.两个线程调用了相同数量的next；2.线程1调用next(a)并进入C代码中的 PyIter_Next 调用。在next（gen）的第一条字节码上，CPython决定切换线程；3.线程2调用next(b)，由于仍需要新的元素，C代码继续执行PyIter_Next调用；4.线程2完成对PyIter_Next的调用并返回一个元素。在某个时刻，CPython再次决定切换线程；5.线程1恢复运行，完成对PyIter_Next的调用，然后运行下面两行： tdo->numread++;
tdo->values[i] = value;
但是线程2已经设置了tdo->values[i]！
这已经足够说明tee不是线程安全的，因为我们丢失了线程2放入tdo->values[i]的值。但这并不能解释崩溃的原因。
假设i是56。由于两个线程都调用了tdo->numread++，它现在变成了58，超过了tdo->values分配的大小57。在线程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不是线程安全的。
我们可以在迭代器的__next__中加锁，而是在tee.__next__周围放置一个锁。这意味着每次都将由单个线程调用整个teedataobject.__getitem__方法。我在答案开头给出了一个简短的实现。它是一个可供替换使用的线程安全tee。 唯一它没有实现但tee实现了的是 - pickle。由于锁不可pickle，因此添加它并不是微不足道的事情。但当然，它是可以做到的。
itertools.tee()及其返回的迭代器是使用C代码实现的。这意味着GIL应该保护它免受多个线程同时调用。它可能会正常工作，并且不会崩溃解释器，但不能保证是线程安全的。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)))