Python 中合并多个阻塞生成器函数

假设有如下情况，我们需要合并两个无限流数据（可能会来自阻塞资源，如套接字），实现非确定性地合并数据，即按数据的到达顺序进行合并。例如，假设有两个迭代器 iter1 和 iter2，则期望合并后的结果等价于 merged。

iter1 : 1 2 3     4   5 ...
iter2 :       1 2   3   ... 
merged: 1 2 3 1 2 4 3 5 ...

   --- > increasing time ---> 

为了完成此任务，我们需要并发程序，但 Python 中是否存在更合适的解决方案呢？同时也希望适用于 Python 2.6。

2. 解决方案

1）使用线程和队列

我们可以将套接字迭代移动到后台线程中，并使用队列将每个套接字接收到的数据发送到主线程。然后主线程可以按顺序从队列中获取数据。

以下代码展示了使用多线程和队列的解决方案：

import socket
import time
from Queue import Queue
from threading import Thread

HOST = "127.0.0.1"
PORT = 8008


def iterate_socket(sock):
    while True:
        data = sock.recv(1024)
        yield data
        if not data:  # End of the stream
            return


def consume(q, s):
    for i in s:
        q.put(i)


def merge(xs, ys):
    q = Queue()
    iters = [xs, ys]
    for it in iters:
        t = Thread(target=consume, args=(q, it))
        t.start()

    done = 0
    while True:
        out = q.get()
        if out == '':  # End of the stream.
            done += 1
            if done == len(iters):  # When all iters are done, break out.
                return
        else:
            yield out


sock1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock1.connect((HOST, PORT))
time.sleep(1)
sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock2.connect((HOST, PORT))

iter1 = iterate_socket(sock1)
iter2 = iterate_socket(sock2)

for msg in merge(iter1, iter2):
    print(msg)

2）使用协程

协程是另一种实现非确定性合并的方式。协程允许在多个函数之间切换执行，而不会阻塞。我们可以使用协程从每个迭代器中获取数据，并将它们合并成一个流。

以下代码展示了使用协程的解决方案：

import asyncio

async def iterate_socket(sock):
    while True:
        data = await sock.recv(1024)
        if not data:
            return
        yield data


async def merge(xs, ys):
    iters = [xs, ys]
    while True:
        for it in iters:
            try:
                data = await it
                yield data
            except StopAsyncIteration:
                pass


async def main():
    sock1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock1.connect((HOST, PORT))
    await asyncio.sleep(1)
    sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock2.connect((HOST, PORT))

    iter1 = iterate_socket(sock1)
    iter2 = iterate_socket(sock2)

    for msg in merge(iter1, iter2):
        print(msg)


asyncio.run(main())

3）使用生成器函数

我们可以使用生成器函数来实现非确定性合并。生成器函数允许我们暂停函数的执行并返回一个值，然后再恢复执行并继续从该值开始。我们可以使用生成器函数从每个迭代器中获取数据，并将它们合并成一个流。

以下代码展示了使用生成器函数的解决方案：

def iterate_socket(sock):
    while True:
        data = sock.recv(1024)
        if not data:
            return
        yield data


def merge(xs, ys):
    iters = [xs, ys]
    while True:
        for it in iters:
            try:
                data = next(it)
                yield data
            except StopIteration:
                pass


sock1 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock1.connect((HOST, PORT))
time.sleep(1)
sock2 = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock2.connect((HOST, PORT))

iter1 = iterate_socket(sock1)
iter2 = iterate_socket(sock2)

for msg in merge(iter1, iter2):
    print(msg)