Python搭建WebSocket服务详解

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

使用Python的websockets库构建WebSocket服务是高效且直观的方案，1. 因其基于asyncio，天然支持高并发异步I/O，每个连接由独立协程处理，通过async for循环接收消息，利用asyncio.gather实现高效广播；2. 服务器通过websockets.serve启动，客户端用websockets.connect连接，代码简洁清晰；3. 常见问题如死连接可通过设置ping_interval和ping_timeout启用心跳机制解决；4. 错误处理需捕获ConnectionClosed异常并在finally中清理连接；5. 安全性方面应使用WSS加密通信并校验Origin头防止跨站劫持；6. 优化策略包括采用Protobuf等高效序列化格式、避免阻塞操作、结合Redis Pub/Sub实现分布式消息同步，最终构建稳定、可扩展的WebSocket服务。

Python构建WebSocket服务，用websockets库是目前我个人觉得最直接、高效的方案。它基于asyncio，天生就是为处理大量并发连接而生，代码写起来也清晰。核心就是创建一个异步服务器，然后处理进出的消息。这套方案不仅性能好，而且开发体验也相当流畅，省去了很多底层协议处理的麻烦。

解决方案

要构建一个基本的WebSocket服务，你需要一个服务器端和一个客户端。服务器端通常使用websockets库的serve函数来启动，并定义一个异步处理函数来响应客户端消息。

服务器端示例 (server.py):

import asyncio
import websockets
import json # 假设我们可能需要处理JSON数据

# 维护所有连接的客户端，用于广播或其他管理
connected_clients = set()

async def handler(websocket, path):
    """
    处理每个新的WebSocket连接。
    `websocket` 是当前连接的实例。
    `path` 是客户端请求的URL路径（通常不怎么用，除非你根据路径区分服务）。
    """
    print(f"客户端 {websocket.remote_address} 已连接。")
    connected_clients.add(websocket) # 将新连接加入集合

    try:
        async for message in websocket:
            # 收到消息，可以做一些处理，比如解析JSON
            print(f"收到来自 {websocket.remote_address} 的消息: {message}")
            try:
                data = json.loads(message)
                if data.get("type") == "broadcast":
                    response_message = f"广播消息: {data.get('content', '无内容')}"
                    # 广播给所有在线客户端
                    await asyncio.gather(*[client.send(response_message) for client in connected_clients if client != websocket])
                else:
                    await websocket.send(f"服务器回应: {message}") # 回复给发送者
            except json.JSONDecodeError:
                await websocket.send(f"服务器回应 (非JSON): {message}")
    except websockets.exceptions.ConnectionClosedOK:
        print(f"客户端 {websocket.remote_address} 正常关闭连接。")
    except websockets.exceptions.ConnectionClosedError as e:
        print(f"客户端 {websocket.remote_address} 连接异常关闭: {e}")
    finally:
        connected_clients.remove(websocket) # 移除断开的连接
        print(f"客户端 {websocket.remote_address} 已断开。")

async def main():
    # 启动WebSocket服务器
    # `websockets.serve` 是一个异步上下文管理器
    async with websockets.serve(handler, "localhost", 8765):
        print("WebSocket服务器已启动，监听 ws://localhost:8765")
        await asyncio.Future()  # 保持服务器运行直到被手动停止

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

客户端示例 (client.py):

import asyncio
import websockets
import json

async def send_message():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        print("已连接到服务器。")

        # 发送一条普通消息
        await websocket.send("Hello from client!")
        response = await websocket.recv()
        print(f"收到服务器回应: {response}")

        # 发送一条广播消息
        broadcast_data = {"type": "broadcast", "content": "大家注意，这是一条广播！"}
        await websocket.send(json.dumps(broadcast_data))
        # 广播消息通常不会有直接回复，除非服务器特意发送

        # 保持连接，等待其他广播消息
        print("等待其他客户端的广播消息 (按 Ctrl+C 退出)...")
        try:
            while True:
                incoming_message = await websocket.recv()
                print(f"收到广播: {incoming_message}")
        except asyncio.CancelledError:
            print("客户端已关闭。")
        except websockets.exceptions.ConnectionClosedOK:
            print("服务器已关闭连接。")

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

运行步骤：

1. 先运行 python server.py 启动服务器。
2. 再运行 python client.py 启动客户端，你可以多开几个客户端实例来测试广播功能。

为什么选择Python的websockets库来构建服务？

在Python的异步生态里，websockets库可以说是一个非常成熟且符合直觉的选择。我个人觉得，它最吸引人的地方在于它完美地融入了asyncio的协程模型。你想想看，WebSocket本身就是一种持久连接、双向通信的协议，天然地需要处理大量的并发I/O。如果用传统的阻塞式I/O，或者自己去手动管理线程池、进程池，那简直是噩梦。asyncio的出现，让Python在处理高并发网络应用上有了质的飞跃，而websockets库正是站在这个巨人的肩膀上。

它抽象掉了WebSocket协议的各种底层细节，比如握手、帧的打包与解包、心跳维持等等。这意味着开发者可以把精力完全放在业务逻辑上，而不是被协议的实现细节所困扰。我记得早期尝试用Python做一些TCP层面的通信时，光是协议解析就让人头大，更别提WebSocket这种更复杂的了。有了websockets，这些都变得像呼吸一样自然。

另外，它对RFC 6455（WebSocket协议标准）的遵循度很高，这保证了你构建的服务能与各种主流的WebSocket客户端（浏览器、移动应用、其他编程语言的客户端）良好兼容。一个符合标准的库，总是能让人用得更安心。而且，社区活跃，遇到问题也比较容易找到解决方案或者得到帮助，这在实际项目中是相当重要的。

如何在WebSocket服务中处理并发连接和消息？

处理并发连接和消息是WebSocket服务的核心挑战之一，但websockets库结合asyncio让这变得相对简单。其秘密武器在于async/await语法和事件循环。

每个新的WebSocket连接，在websockets服务器端都会被websockets.serve函数分配一个独立的协程来处理（就是上面示例中的handler函数）。这个协程会一直运行，直到客户端断开连接。这意味着，即使有成千上万个客户端同时连接，每个客户端都有自己的“专属”协程在后台默默工作，互不干扰。当一个协程在等待I/O（比如等待客户端发送消息await websocket.recv()）时，它会主动“让出”CPU，让事件循环去处理其他已经准备好的任务（比如其他客户端发来的消息，或者新的连接请求）。这种非阻塞的I/O模式，就是实现高并发的关键。

至于消息处理，通常的做法是在handler协程内部使用一个async for message in websocket:循环。这个循环会持续监听当前连接的所有入站消息。每收到一条消息，你就可以在循环体内部进行处理。如果需要向特定客户端发送消息，就用await websocket.send(message)；如果需要向所有在线客户端广播消息，就像示例中那样，你需要维护一个所有活跃连接的集合（比如connected_clients），然后遍历这个集合，对每个客户端调用await client.send(message)。

这里有个小技巧：当需要同时向多个客户端发送消息时，使用asyncio.gather(*[client.send(...) for client in connected_clients])会比简单的循环更高效，因为它会并发地执行所有发送操作，而不是一个接一个地阻塞等待。当然，实际项目中，你可能还需要考虑消息队列（如Redis Pub/Sub）来在多个WebSocket服务实例之间同步消息，尤其是在构建分布式系统时。

WebSocket通信中的常见问题与优化策略有哪些？

在实际部署和运行WebSocket服务时，确实会遇到一些问题，也有不少可以优化的点。

一个很常见的问题是“死连接”。客户端可能在没有正常关闭连接的情况下断开（比如网络突然中断、浏览器崩溃），服务器端并不会立即知道。这时，你的connected_clients集合里可能就包含了这些无效的连接，向它们发送消息会失败。解决办法通常是使用WebSocket的心跳机制（Ping/Pong帧）。websockets库内置了对心跳的支持，你可以在serve函数中设置ping_interval和ping_timeout参数。比如，ping_interval=20表示每20秒发送一个Ping帧给客户端，如果客户端在ping_timeout时间内没有回复Pong，服务器就会认为连接已死并自动关闭。这能有效清理无效连接。

另一个是错误处理。网络通信总是充满不确定性。客户端连接可能突然中断，发送的数据格式可能不符合预期。在handler协程中，使用try...except websockets.exceptions.ConnectionClosedOK和websockets.exceptions.ConnectionClosedError来捕获连接关闭的各种情况，并在finally块中清理资源（比如从connected_clients集合中移除连接），这是非常重要的。对于业务逻辑内部的错误，比如JSON解析失败，也应该有相应的try...except块来防止服务崩溃，并向客户端返回有意义的错误信息。

安全性也是不可忽视的一环。首先，生产环境务必使用WSS (WebSocket Secure)，即基于TLS/SSL的WebSocket。这和HTTPS是类似的，可以加密通信内容，防止中间人攻击。在websockets.serve中传入ssl_context参数即可。其次，要进行来源验证（Origin Check）。你可能不希望任何网站都能连接到你的WebSocket服务。通过检查请求头中的Origin字段，只允许来自你信任的域名连接，可以有效防止跨站WebSocket劫持（Cross-Site WebSocket Hijacking）。在handler函数内部可以访问websocket.request_headers['Origin']来进行判断。

最后，关于性能优化，除了上面提到的并发处理，如果你预期有海量的连接或消息吞吐，可以考虑：

1. 消息序列化/反序列化：如果传输的是结构化数据，选择高效的序列化协议，如Protobuf或MessagePack，它们通常比JSON更紧凑，解析速度更快。
2. 避免不必要的计算：在handler中，尽量减少同步阻塞的计算，将耗时操作放到线程池或单独的进程中处理，或者使用asyncio.to_thread()。
3. 负载均衡：当单机性能达到瓶颈时，将WebSocket服务部署到多台服务器上，并通过负载均衡器（如Nginx）进行流量分发。这通常需要结合外部消息队列（如Redis Pub/Sub、Kafka）来实现跨服务器的消息广播和状态同步。

到这里，我们也就讲完了《Python搭建WebSocket服务详解》的内容了。个人认为，基础知识的学习和巩固，是为了更好的将其运用到项目中，欢迎关注golang学习网公众号，带你了解更多关于Python,WebSockets,websocket,并发,asyncio的知识点！