Python asyncio 缓冲和处理数据

Question

我在使用asyncio事件循环时遇到了一些关于CPU占用高的任务的问题。这些问题主要出现在处理接收到的数据缓冲区和从中构建数据包的时候。我尝试过使用执行器来处理这些CPU密集型的任务，但在从缓冲区中移除数据包时，保持缓冲区的顺序就变得很麻烦。

我希望找到一种最佳实践的方法，来实现以下功能，而不让CPU占用高的任务在事件循环中执行。

import asyncio
import struct

class Reader(asyncio.Protocol):
    def __init__(self):
        self.extra = bytearray()

    def data_received(self, data):
        self.extra.extend(data)
        packet = get_packet(bytes(self.extra))
        if packet:
            del self.extra[:len(packet)]
            if verify_hash(packet):  # CPU intensive
                asyncio.async(distribute(packet))  # Some asyncio fan-out callback


def get_packet(data):  # CPU intensive
    if len(data) > HEADER_SIZE:
        payload_size, = struct.unpack_from(HEADER_FORMAT, data)
        if len(data) >= HEADER_SIZE + payload_size:
            return data[:HEADER_SIZE + payload_size]
    return None

loop = asyncio.get_event_loop()
loop.run_until_complete(loop.create_server(Reader, '0.0.0.0', 8000))
loop.run_forever()

Answer 1

我会尝试把整个数据包处理的逻辑整理一下，把繁重的任务拆分成几个小部分。比如可以使用MD5哈希：

@asyncio.coroutine
def verify_hash(packet):
    m = hashlib.md5()
    for i in range(len(packet) // 4096 + 1):
        yield m.update(packet[i:i+4096])
    return m.digest() == signature


@asyncio.coroutine
def handle_packet(packet):
    verified = yield from verify_hash(packet)
    if verified:
        yield from distribute(packet)


class Reader(asyncio.Protocol):
    def __init__(self):
        self.extra = bytearray()

    def data_received(self, data):
        self.extra.extend(data)
        packet = get_packet(bytes(self.extra))
        if packet:
            del self.extra[:len(packet)]
            asyncio.async(handle_packet(packet))

要注意，数据包的到达速度可能会比Reader处理的速度快得多，所以要确保监控系统的负载，并在需要的时候停止接收数据包。不过这又是另一个话题了 :)

Answer 2

你想尽可能快地处理进入Reader的所有数据，但不能让多个线程或进程同时处理这些数据；这就是你之前在使用执行器时遇到竞争条件的原因。相反，你应该启动一个工作进程，这个进程可以逐个处理所有的数据包，使用multiprocessing.Queue来将数据从父进程传递给工作进程。然后，当工作进程构建、验证并准备好分发一个有效的数据包时，它会使用另一个multiprocessing.Queue将数据包发送回父进程中的一个线程，这个线程可以使用线程安全的call_soon_threadsafe方法来安排distribute的执行。

下面是一个未经测试的示例，应该能给你一个如何实现这个的概念：

import asyncio
import struct
from concurrent.futures.ProcessPoolExecutor
import threading


def handle_result_packets():
    """ A function for handling packets to be distributed.

    This function runs in a worker thread in the main process.

    """
    while True:
        packet = result_queue.get()
        loop.call_soon_threadsafe(asyncio.async, distribute(packet))

def get_packet():  # CPU intensive
    """ Handles processing all incoming packet data.

    This function runs in a separate process.

    """
    extra = bytearray()
    while True:
        data = data_queue.get()
        extra.extend(data)
        if len(data) > HEADER_SIZE:
            payload_size, = struct.unpack_from(HEADER_FORMAT, data)
            if len(data) >= HEADER_SIZE + payload_size:
                packet = data[:HEADER_SIZE + payload_size]
                del extra[:len(packet)]
                if verify_hash(packet):
                    result_queue.put(packet)


class Reader(asyncio.Protocol):
    def __init__(self):
        self.extra = bytearray()
        self.t = threading.Thread(target=handle_result_packets)
        self.t.start()

    def data_received(self, data):
        data_queue.put(data)


if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    data_queue = multiprocessing.Queue()
    result_queue = multiprocessing.Queue()
    p = multiprocessing.Process(target=get_packet)
    p.start()
    loop.run_until_complete(loop.create_server(Reader, '0.0.0.0', 8000))
    loop.run_forever()