一、技术选型对比

二、核心系统实现

1. WebSocket服务端

import asyncio
import websockets
from collections import defaultdict

class ChatServer:
    def __init__(self):
        self.rooms = defaultdict(set)
        self.user_ws = {}
        self.user_room = {}

    async def handle_connection(self, websocket, path):
        try:
            user_id = await websocket.recv()
            room_id = path.strip('/')
            
            self.user_ws[user_id] = websocket
            self.user_room[user_id] = room_id
            self.rooms[room_id].add(user_id)
            
            await self.broadcast(f"User {user_id} joined", room_id)
            
            async for message in websocket:
                await self.broadcast(f"{user_id}: {message}", room_id)
                
        except websockets.ConnectionClosed:
            await self.handle_disconnect(user_id)

    async def handle_disconnect(self, user_id):
        room_id = self.user_room.get(user_id)
        if room_id:
            self.rooms[room_id].discard(user_id)
            await self.broadcast(f"User {user_id} left", room_id)
        self.user_ws.pop(user_id, None)
        self.user_room.pop(user_id, None)

    async def broadcast(self, message, room_id):
        if room_id in self.rooms:
            tasks = [
                ws.send(message)
                for user_id in self.rooms[room_id]
                if (ws := self.user_ws.get(user_id))
            ]
            await asyncio.gather(*tasks)

2. 客户端连接管理

async def chat_client():
    uri = "ws://localhost:8765/room1"
    async with websockets.connect(uri) as websocket:
        # 身份认证
        await websocket.send("user123")
        
        # 接收消息任务
        receive_task = asyncio.create_task(
            receive_messages(websocket)
        )
        
        # 发送消息
        try:
            while True:
                message = input("> ")
                await websocket.send(message)
        finally:
            receive_task.cancel()

async def receive_messages(websocket):
    async for message in websocket:
        print(f"n< {message}")

3. 性能优化配置

# 启动优化后的WebSocket服务
start_server = websockets.serve(
    ChatServer().handle_connection,
    "0.0.0.0",
    8765,
    max_size=2**20,  # 1MB消息限制
    ping_interval=30,
    ping_timeout=90,
    compression="deflate"
)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

三、高级功能扩展

1. 消息持久化存储

class MessageLogger:
    def __init__(self):
        self.redis = redis.Redis()
        
    async def log_message(self, room_id, user_id, message):
        msg_data = {
            'timestamp': time.time(),
            'user': user_id,
            'content': message
        }
        self.redis.lpush(
            f"chat:{room_id}",
            json.dumps(msg_data)
        )
        
    async def get_history(self, room_id, limit=100):
        return [
            json.loads(msg)
            for msg in self.redis.lrange(
                f"chat:{room_id}", 0, limit-1)
        ]

2. 分布式扩展方案

class PubSubManager:
    def __init__(self):
        self.pubsub = redis.Redis().pubsub()
        
    async def subscribe(self, room_id):
        self.pubsub.subscribe(f"chat:{room_id}")
        
    async def publish(self, room_id, message):
        self.redis.publish(f"chat:{room_id}", message)
        
    async def listen(self):
        async for message in self.pubsub.listen():
            if message['type'] == 'message':
                yield message

四、性能测试数据