如何将asyncio StreamReader广播给几个消费者？

Question

我正在尝试使用aiohttp来制作一种高级反向代理

我想获取HTTP请求的内容，并将其传递给新的HTTP请求，而无需将其拉入内存。虽然只有上游，但任务相当简单：aiohttp服务器以StreamReader的形式返回请求内容，而aiohttp客户端可以接受StreamReader作为请求主体

问题是，我想将原始请求发送到几个上游，或者，例如，同时将内容发送到上游并将其写入磁盘

是否有一些工具可以播放StreamReader的内容

我曾尝试制作一些天真的广播，但在大型对象上失败了。我做错了什么

class StreamBroadcast:
    async def __do_broadcast(self):
        while True:
            chunk = await self.__source.read(self.__n)
            if not chunk:
                break
            for output in self.__sinks:
                output.feed_data(chunk)
        for output in self.__sinks:
            output.feed_eof()

    def __init__(self, source: StreamReader, sinks_count: int, n: int = -1):
        self.__source = source
        self.__n = n
        self.__sinks = [StreamReader() for i in range(sinks_count)]
        self.__task = asyncio.create_task(self.__do_broadcast())

    @property
    def sinks(self) -> Iterable[StreamReader]:
        return self.__sinks

    @property
    def ready(self) -> Task:
        return self.__task

Answer 1

嗯，我查看了asyncio源代码，发现我应该使用Transport在流上泵送数据。这是我的解决办法

import asyncio
from asyncio import StreamReader, StreamWriter, ReadTransport, StreamReaderProtocol
from typing import Iterable


class _BroadcastReadTransport(ReadTransport):
    """
    Internal class, is not meant to be instantiated manually
    """

    def __init__(self, source: StreamReader, sinks: Iterable[StreamReader]):
        super().__init__()
        self.__source = source
        self.__sinks = tuple(StreamReaderProtocol(s) for s in sinks)
        for sink in sinks:
            sink.set_transport(self)
        self.__waiting_for_data = len(self.__sinks)

        asyncio.create_task(self.__broadcast_next_chunk(), name='initial-chunk-broadcast')

    def is_reading(self):
        return self.__waiting_for_data == len(self.__sinks)

    def pause_reading(self):
        self.__waiting_for_data -= 1

    async def __broadcast_next_chunk(self):
        data = await self.__source.read()
        if data:
            for sink in self.__sinks:
                sink.data_received(data)
            if self.is_reading():
                asyncio.create_task(self.__broadcast_next_chunk())
        else:
            for sink in self.__sinks:
                sink.eof_received()

    def resume_reading(self):
        self.__waiting_for_data += 1
        if self.__waiting_for_data == len(self.__sinks):
            asyncio.create_task(self.__broadcast_next_chunk(), name='chunk-broadcast')

    @property
    def is_completed(self):
        return self.__source.at_eof()


class StreamBroadcast:
    def __init__(self, source: StreamReader, sinks_count: int):
        self.__source = source
        self.__sinks = tuple(StreamReader() for _ in range(sinks_count))
        self.__transport = _BroadcastReadTransport(self.__source, self.__sinks)

    @property
    def sinks(self) -> Iterable[StreamReader]:
        return self.__sinks

    @property
    def is_completed(self):
        return self.__transport.is_completed

希望有一次我会把它打包到pip模块