码迷,mamicode.com
首页 > 编程语言 > 详细

一日一技:如何让 itertools.tee 线程安全

时间:2020-12-08 11:59:32      阅读:6      评论:0      收藏:0      [点我收藏+]

标签:image   效果   info   end   理解   生成   mod   参考资料   lock   

一日一技:如何让 itertools.tee 线程安全

技术图片

摄影:产品经理
一起吃:kingname & 产品经理
在上一篇文章中,我们说到了,itertools.tee不是线程安全的,并给出了一个例子,如下图所示:
技术图片

在两个线程里面同时运行分裂出来的生成器对象,就会导致报错。

现在,你想看看itertools.tee的源代码,但是你会发现,在 PyCharm 里面,它的源代码如下图所示:

技术图片
这是因为,在 CPython 中,itertools.tee底层是通过 C 语言实现的,所以你不能在 PyCharm 中看到它的源代码。但是你可以通过阅读 Python 的源代码中的 Modules/itertoolsmodule.c 文件[1],找到它的实现算法。

导致问题的核心部分在如下图所示的两段代码中:
技术图片
技术图片

大家看不懂也没有关系,根据我上一篇文章中使用 Python 实现的简化版本就足够帮助理解了。

我们使用简化版本来解释其中线程不安全的地方:


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)

当两个线程同时运行到if not queue时,发现当前各自的队列都是空的,于是进入value = next(g)获取下一个值。其中,线程 A 先进入那么几毫秒。然后线程 B 进入value = next(g)。但由于此时线程 A 中的next(g)正在运行,尚未结束,线程 B 又跑来运行,于是就导致了报错的发生。Python 中,生成器不是线程安全的。

那么如何让itertools.tee分裂出来的多个生成器可以在多线程中运行呢?其关键因素就是让value = next(g)这一行一次只能让一个线程运行。所以我们可以通过加锁来实现。


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))

我们来看看运行效果:
技术图片

多线程完美运行。

参考资料

[1]
Python 的源代码中的Modules/itertoolsmodule.c文件: https://github.com/python/cpython/blob/master/Modules/itertoolsmodule.c

技术图片

一日一技:如何让 itertools.tee 线程安全

标签:image   效果   info   end   理解   生成   mod   参考资料   lock   

原文地址:https://blog.51cto.com/15023263/2558921

(0)
(0)
   
举报
评论 一句话评论(0
登录后才能评论!
© 2014 mamicode.com 版权所有  联系我们:gaon5@hotmail.com
迷上了代码!