Python 多个异步循环使用同一WebSocket连接的问题咨询

核心逻辑是在全局/外层统一创建唯一的WebSocket连接实例，再将该实例传入两个异步任务共同使用，websockets库原生支持单连接下的并发读写操作，一个协程负责收消息、另一个负责发消息不会出现冲突。

修改后的可运行代码 #

# Imports
import asyncio
import websockets
from datetime import datetime

websocket_url = "wss://localhost:5000/"

# 模拟拉取数据
async def fetch_data():
    print("Fetching started")
    await asyncio.sleep(2)
    return {"data": 2}

# 接收并处理WebSocket消息，接收外部传入的连接实例
async def retrieve_websocket(ws):
    while True:
        msg = await ws.recv()
        print(msg)
        # 满足条件时执行对应操作

# 定时拉取数据并发送WebSocket消息，接收外部传入的连接实例
async def periodic_fetch(ws):
    while True:
        print(datetime.now())
        fetch_task = asyncio.create_task(fetch_data())
        wait_task = asyncio.create_task(asyncio.sleep(60))

        res = await fetch_task
        # 发送消息到WebSocket
        await ws.send("Websocket message")
        # 等待剩余的定时时长
        await wait_task

# 主入口逻辑：统一创建连接，再启动两个异步任务
async def main():
    async with websockets.connect(websocket_url) as ws:
        # 同时启动两个协程，传入同一个连接实例
        await asyncio.gather(
            retrieve_websocket(ws),
            periodic_fetch(ws)
        )

if __name__ == "__main__":
    asyncio.run(main())

注意事项 #

• 如果需要处理连接断开重连的场景，可以在main函数外层加异常捕获和重试逻辑，连接断开后重新创建连接再启动两个任务即可
• 如果你有多个协程需要同时发送消息，建议加一个异步队列缓冲发送请求，避免并发send的边界冲突，单发送协程+单接收协程的场景不需要额外加锁

内容的提问来源于stack exchange，提问作者HaakonFlaar