使用 asyncio 和 websockets 实现实时视频流推送

本文将指导开发者如何使用 asyncio 和 websockets 库，在 Python 中构建一个能够实时广播视频流预测结果的服务器。关键在于使用 asyncio.wait 替代 websockets.broadcast，以解决高并发场景下客户端无法接收数据的问题。我们将深入探讨这两种方法之间的差异，并提供完整的代码示例和注意事项，帮助读者理解并实现高效的实时数据广播。

问题背景

在构建实时视频流处理应用时，通常需要将视频帧进行预测，并将预测结果广播给多个客户端。使用 websockets 库可以方便地实现 WebSocket 通信，但当在高并发场景下使用 websockets.broadcast 时，可能会遇到客户端无法及时接收数据的问题。

解决方案：使用 asyncio.wait 替代 websockets.broadcast

原来的代码中使用 websockets.broadcast(clients, result) 来广播预测结果，这在某些情况下会导致阻塞，特别是当客户端数量较多或者预测过程耗时较长时。解决方案是将 websockets.broadcast 替换为 asyncio.wait([ws.send(result) for ws in clients])。

以下是修改后的服务器端代码：

import websockets
import cv2
import asyncio
import time

def predict(image):
    # 模拟预测函数
    time.sleep(0.1) # 模拟预测耗时
    return "test"

async def echo(websocket, path):
    global vidCap, i
    while True:
        ret, image = vidCap.read()
        if ret:
            start = time.time()
            result = predict(image)
            # 使用 asyncio.wait 替代 websockets.broadcast
            await asyncio.wait([ws.send(result) for ws in clients])
            end = time.time()
            print("exec time：%f s" % (end - start))

async def handler(websocket, path):
    clients.add(websocket)
    try:
        await echo(websocket, path)
    finally:
        clients.remove(websocket)

async def serve():
    start_server = await websockets.serve(handler, "localhost", 8765)
    await start_server.wait_closed()

if __name__ == '__main__':
    vidCap = cv2.VideoCapture(0) # 使用摄像头代替视频文件
    clients = set()
    asyncio.run(serve())

代码解释

1. asyncio.wait 的作用: asyncio.wait 接收一个 awaitable 对象的集合，并并发地执行它们。它会等待所有 awaitable 对象完成，或者直到满足指定的条件（例如，超时）。
2. [ws.send(result) for ws in clients]: 这是一个列表推导式，它为每个客户端 ws 创建一个 ws.send(result) 的 awaitable 对象。ws.send(result) 是一个异步函数，用于向客户端发送数据。
3. 替换的原因: websockets.broadcast 内部可能没有充分利用 asyncio 的并发能力，导致在高并发场景下阻塞。而 asyncio.wait 能够更有效地并发执行多个 send 操作，从而提高性能。

客户端代码

客户端代码无需修改，如下所示：

import websockets
import asyncio
import time

async def get_result(uri):
    async with websockets.connect(uri) as websocket:
        while(True):
            try:
                start = time.time()
                recv_text = await websocket.recv()
                print(recv_text)
                end = time.time()
                print("exec：%f s" % (end - start))
            except:
                pass

if __name__ == '__main__':
    asyncio.run(get_result("ws://127.0.0.1:8765/ws"))

websockets.broadcast vs asyncio.wait

• websockets.broadcast: 这是一个便捷的函数，用于向所有连接的客户端广播消息。但是，它可能不是最有效的并发方法，尤其是在客户端数量很多的情况下。
• asyncio.wait: 提供了更细粒度的控制，允许并发地执行多个 send 操作。这可以显著提高性能，并避免阻塞。

注意事项

• 错误处理: 在实际应用中，应该添加适当的错误处理机制，例如，捕获 websocket.send 可能抛出的异常。
• 资源管理: 确保正确关闭 WebSocket 连接，以避免资源泄漏。
• 性能优化: 根据实际情况调整预测函数的性能，以及 WebSocket 连接的参数，以获得最佳性能。
• 数据格式: 确保发送的数据格式是客户端能够正确解析的。

总结

通过使用 asyncio.wait 替代 websockets.broadcast，可以显著提高实时视频流预测结果广播的性能。这种方法能够更有效地利用 asyncio 的并发能力，避免阻塞，并提高系统的响应速度。在实际应用中，应该根据具体情况进行调整和优化，以获得最佳性能。