asyncio.wait优化WebSocket广播性能

在文章实战开发的过程中，我们经常会遇到一些这样那样的问题，然后要卡好半天，等问题解决了才发现原来一些细节知识点还是没有掌握好。今天golang学习网就整理分享《asyncio.wait 优化 WebSocket 广播方法》，聊聊，希望可以帮助到正在努力赚钱的你。

本文将探讨在使用 websockets 库进行 WebSocket 广播时，如何避免程序阻塞，确保客户端能够及时接收消息。起初，使用 websockets.broadcast() 在 while True 循环中广播预测结果时，发现客户端无法正常接收数据。经过分析和实验，通过将 websockets.broadcast() 替换为 asyncio.wait([ws.send(result) for ws in clients])，问题得到了有效解决。本文将详细介绍这种方法，并解释其背后的原因。

问题描述

在使用 websockets 库构建一个 WebSocket 服务器，用于实时传输视频帧的预测结果时，遇到了一个问题：服务器在 while True 循环中使用 websockets.broadcast() 向所有客户端广播预测结果，但客户端无法及时接收到数据，甚至只有在服务器停止运行时才能接收到数据。

原始代码片段如下：

import websockets
import cv2
import asyncio
import time

def predict(image):
    return "test"

async def echo(websocket, path):
    global vidCap,i
    while True:
        ret, image = vidCap.read()
        if ret:
            start = time.time()
            result=predict(image)
            websockets.broadcast(clients,result)
            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('rtsp://xxx.xxx.xx') #rtsp or video
    clients = set()
    asyncio.run(serve())

解决方案

经过研究，发现 websockets.broadcast() 可能会阻塞事件循环，导致其他任务无法执行，从而影响客户端接收消息。为了解决这个问题，将 websockets.broadcast() 替换为 asyncio.wait([ws.send(result) for ws in clients])。

修改后的代码片段如下：

import websockets
import cv2
import asyncio
import time

def predict(image):
    return "test"

async def echo(websocket, path):
    global vidCap,i
    while True:
        ret, image = vidCap.read()
        if ret:
            start = time.time()
            result=predict(image)
            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('rtsp://xxx.xxx.xx') #rtsp or video
    clients = set()
    asyncio.run(serve())

原理解释

websockets.broadcast(clients, result) 是一个同步操作，它会依次向 clients 集合中的每个 WebSocket 连接发送 result。如果某个客户端连接速度较慢，或者网络状况不佳，websockets.broadcast() 可能会阻塞，导致整个事件循环暂停，影响其他客户端的消息发送。

asyncio.wait([ws.send(result) for ws in clients]) 则使用 asyncio.wait() 并发地执行 ws.send(result) 协程。asyncio.wait() 会等待所有协程完成，但它不会阻塞事件循环。这意味着，即使某个客户端连接速度较慢，其他客户端仍然可以接收到消息。

具体来说，[ws.send(result) for ws in clients] 创建了一个包含所有 ws.send(result) 协程的列表。asyncio.wait() 接收这个列表，并并发地执行这些协程。由于 ws.send() 是一个异步操作，它不会阻塞事件循环。asyncio.wait() 会在所有 ws.send() 协程完成后返回，但它不会阻止其他任务的执行。

注意事项

• 确保所有客户端连接都已建立，并且客户端的 websocket 对象已经添加到 clients 集合中。
• 在处理视频帧时，尽量减少 predict() 函数的执行时间，以提高服务器的吞吐量。可以使用异步任务来执行预测，避免阻塞事件循环。
• 根据实际情况调整视频帧的传输速率，避免客户端接收速度过慢导致消息堆积。
• 错误处理：在 ws.send() 中添加适当的异常处理，例如 try...except 块，以防止单个客户端连接的问题影响到整个广播过程。

总结

通过将 websockets.broadcast() 替换为 asyncio.wait([ws.send(result) for ws in clients])，可以实现非阻塞的 WebSocket 广播，确保服务器能够持续处理视频帧并及时将预测结果发送给所有客户端。这种方法可以有效解决在高并发场景下，由于单个客户端连接速度较慢而导致的消息阻塞问题。在实际应用中，应根据具体情况进行调整和优化，以达到最佳的性能。

终于介绍完啦！小伙伴们，这篇关于《asyncio.wait优化WebSocket广播性能》的介绍应该让你收获多多了吧！欢迎大家收藏或分享给更多需要学习的朋友吧~golang学习网公众号也会发布文章相关知识，快来关注吧！