独立消费 websocket 消息（Python）

Question

有没有办法使用两个单独的async for 循环从 websocket 连接中消费两次消息？

运行下面的代码会给出 RuntimeError: cannot call recv 而另一个协程已经在等待下一条消息。

import websockets
import asyncio

async def foo(ws):
    async for msg in ws:
        print(f"foo: {msg}")


async def bar(ws):
    async for msg in ws:
        print(f"bar: {msg}")


async def main():
    async with websockets.connect("wss://echo.websocket.org") as ws:
        asyncio.create_task(foo(ws))
        asyncio.create_task(bar(ws))
        await ws.send("Hello")


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())

Answer 1

您可以创建一个广播函数，该函数读取一次ws，并将每条消息传输到多个生成器。例如（未经测试）：

def broadcast(stream, num):
    # iterate over the stream just once, but put each message into
    # multiple queues
    queues = []
    async def consume():
        async for msg in stream:
            for queue in queues:
                await queue.put(msg)
        for queue in queues:
            await queue.put(None)
    asyncio.create_task(consume())

    # create the queues and return the generators that transmit
    # their contents
    async def transmit(queue):
        while True:
            msg = await queue.get()
            if msg is None:
                break
            yield msg

    iters = []
    for _ in range(num):
        queue = asyncio.Queue()
        iters.append(transmit(queue))
        queues.append(queue)

    return iters

有了这个，你main() 可能看起来像这样：

async def main():
    async with websockets.connect("wss://echo.websocket.org") as ws:
        foo_stream, bar_stream = broadcast(ws, 2)
        asyncio.create_task(foo(foo_stream))
        asyncio.create_task(bar(bar_stream))
        await ws.send("Hello")