怎么使itertools.tee線程安全

 

在兩個線程里面同時運行分裂出來的生成器對象，就會導(dǎo)致報錯。

現(xiàn)在，你想看看itertools.tee的源代碼，但是你會發(fā)現(xiàn)，在 PyCharm 里面，它的源代碼如下圖所示：

 

這是因為，在 CPython 中，itertools.tee底層是通過 C 語言實現(xiàn)的，所以你不能在 PyCharm 中看到它的源代碼。但是你可以通過閱讀 Python 的源代碼中的 Modules/itertoolsmodule.c 文件^[1]，找到它的實現(xiàn)算法。

導(dǎo)致問題的核心部分在如下圖所示的兩段代碼中：

 

 

大家看不懂也沒有關(guān)系，根據(jù)我上一篇文章中使用 Python 實現(xiàn)的簡化版本就足夠幫助理解了。

我們使用簡化版本來解釋其中線程不安全的地方：

def generator():
    for i in range(3):
        yield f'我是你第{i}個爺爺'

def split(g):
    value_list_1 = []
    value_list_2 = []
    def wrap(queue):
        while True:
            if not queue:
                try:
                    value = next(g)
                except StopIteration:
                    return
                value_list_1.append(value)
                value_list_2.append(value)
            yield queue.pop(0)
    g_1 = wrap(value_list_1)
    g_2 = wrap(value_list_2)
    return g_1, g_2

g = generator()
g_1, g_2 = split(g)
for value in g_1:
    print(value)

for value in g_2:
    print(value)

當(dāng)兩個線程同時運行到if not queue時，發(fā)現(xiàn)當(dāng)前各自的隊列都是空的，于是進入value = next(g)獲取下一個值。其中，線程 A 先進入那么幾毫秒。然后線程 B 進入value = next(g)。但由于此時線程 A 中的next(g)正在運行，尚未結(jié)束，線程 B 又跑來運行，于是就導(dǎo)致了報錯的發(fā)生。Python 中，生成器不是線程安全的。

那么如何讓itertools.tee分裂出來的多個生成器可以在多線程中運行呢？其關(guān)鍵因素就是讓value = next(g)這一行一次只能讓一個線程運行。所以我們可以通過加鎖來實現(xiàn)。

import itertools
from threading import Lock


class KingnameTee:
    def __init__(self, tee_obj, lock):
        self.tee_obj = tee_obj
        self.lock = lock

    def __iter__(self):
        return self

    def __next__(self):
        with self.lock:
            return next(self.tee_obj)

    def __copy__(self):
        return KingnameTee(self.tee_obj.__copy__(), self.lock)

def safe_tee(iterable, n=2):
    """tuple of n independent thread-safe iterators"""
    lock = Lock()
    return tuple(KingnameTee(tee_obj, lock) for tee_obj in itertools.tee(iterable, n))

我們來看看運行效果：

 

多線程完美運行。