Python脚本WebSocket客户端如何重连

wen 实用脚本 8

本文目录导读:

Python脚本WebSocket客户端如何重连

  1. 方案一:使用 websockets 库 + 循环重连(推荐)
  2. 方案二:使用 websocket-client 库 + 回调重连
  3. 方案三:带指数退避的智能重连

在 Python WebSocket 客户端中实现自动重连,有多种方法,这里提供几种常用且稳定的实现方案:

使用 websockets 库 + 循环重连(推荐)

import asyncio
import websockets
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class WebSocketClient:
    def __init__(self, uri, max_retries=5, retry_delay=3):
        self.uri = uri
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.websocket = None
        self.running = False
    async def connect(self):
        """建立 WebSocket 连接"""
        try:
            self.websocket = await websockets.connect(self.uri)
            logger.info(f"Connected to {self.uri}")
            return True
        except Exception as e:
            logger.error(f"Failed to connect: {e}")
            return False
    async def reconnect(self):
        """重连逻辑"""
        retries = 0
        while retries < self.max_retries:
            logger.info(f"Reconnecting... attempt {retries + 1}/{self.max_retries}")
            if await self.connect():
                return True
            retries += 1
            if retries < self.max_retries:
                await asyncio.sleep(self.retry_delay)
        logger.error("Max reconnection attempts reached")
        return False
    async def send_message(self, message):
        """发送消息"""
        if self.websocket:
            try:
                await self.websocket.send(message)
                return True
            except websockets.ConnectionClosed:
                logger.warning("Connection closed while sending")
                return False
        return False
    async def receive_messages(self):
        """接收消息的主循环"""
        self.running = True
        while self.running:
            try:
                if not self.websocket:
                    if not await self.reconnect():
                        break
                message = await asyncio.wait_for(
                    self.websocket.recv(), 
                    timeout=30
                )
                self.handle_message(message)
            except asyncio.TimeoutError:
                # 超时处理,保持连接活跃
                try:
                    await self.websocket.ping()
                except:
                    self.websocket = None
                    continue
            except websockets.ConnectionClosed:
                logger.warning("Connection lost, attempting to reconnect...")
                self.websocket = None
                if not await self.reconnect():
                    break
            except Exception as e:
                logger.error(f"Unexpected error: {e}")
                self.websocket = None
                await asyncio.sleep(self.retry_delay)
    def handle_message(self, message):
        """处理接收到的消息"""
        logger.info(f"Received: {message}")
    async def close(self):
        """关闭连接"""
        self.running = False
        if self.websocket:
            await self.websocket.close()
# 使用示例
async def main():
    client = WebSocketClient("wss://echo.websocket.org")
    try:
        # 启动接收消息的任务
        receive_task = asyncio.create_task(client.receive_messages())
        # 等待一段时间后发送消息
        await asyncio.sleep(2)
        await client.send_message("Hello, WebSocket!")
        # 模拟网络断开(仅用于测试)
        await asyncio.sleep(5)
        # 等待接收任务完成
        await receive_task
    except KeyboardInterrupt:
        await client.close()
if __name__ == "__main__":
    asyncio.run(main())

使用 websocket-client 库 + 回调重连

import websocket
import json
import time
import logging
from threading import Thread
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class WebSocketClient:
    def __init__(self, url, reconnect_interval=5, max_reconnect_attempts=10):
        self.url = url
        self.reconnect_interval = reconnect_interval
        self.max_reconnect_attempts = max_reconnect_attempts
        self.ws = None
        self.running = False
        self.reconnect_attempts = 0
    def on_message(self, ws, message):
        """消息回调"""
        logger.info(f"Received: {message}")
    def on_error(self, ws, error):
        """错误回调"""
        logger.error(f"Error: {error}")
    def on_close(self, ws, close_status_code, close_msg):
        """关闭回调"""
        logger.info(f"Connection closed: {close_status_code} - {close_msg}")
        if self.running:
            self.reconnect()
    def on_open(self, ws):
        """打开回调"""
        logger.info("Connection established")
        self.reconnect_attempts = 0
    def connect(self):
        """建立连接"""
        websocket.enableTrace(False)
        self.ws = websocket.WebSocketApp(
            self.url,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close,
            on_open=self.on_open
        )
        # 启动运行线程
        self.running = True
        wst = Thread(target=self.ws.run_forever)
        wst.daemon = True
        wst.start()
    def reconnect(self):
        """重连逻辑"""
        while self.running and self.reconnect_attempts < self.max_reconnect_attempts:
            self.reconnect_attempts += 1
            logger.info(f"Reconnecting... attempt {self.reconnect_attempts}/{self.max_reconnect_attempts}")
            try:
                time.sleep(self.reconnect_interval)
                self.connect()
                break
            except Exception as e:
                logger.error(f"Reconnection failed: {e}")
        if self.reconnect_attempts >= self.max_reconnect_attempts:
            logger.error("Max reconnection attempts reached")
            self.running = False
    def send(self, message):
        """发送消息"""
        if self.ws and self.ws.sock:
            try:
                if isinstance(message, dict):
                    message = json.dumps(message)
                self.ws.send(message)
                logger.info(f"Sent: {message}")
                return True
            except Exception as e:
                logger.error(f"Failed to send message: {e}")
        return False
    def close(self):
        """关闭连接"""
        self.running = False
        if self.ws:
            self.ws.close()
# 使用示例
def main():
    client = WebSocketClient("wss://echo.websocket.org")
    client.connect()
    try:
        # 发送消息
        time.sleep(2)
        client.send("Hello, WebSocket!")
        # 保持运行
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        client.close()
        logger.info("Client stopped")
if __name__ == "__main__":
    main()

带指数退避的智能重连

import asyncio
import websockets
import random
import logging
logger = logging.getLogger(__name__)
class SmartReconnectWebSocket:
    def __init__(self, uri, max_retries=10, base_delay=1, max_delay=60):
        self.uri = uri
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.retry_count = 0
        self.websocket = None
    def get_backoff_time(self):
        """计算指数退避时间"""
        delay = min(self.base_delay * (2 ** self.retry_count), self.max_delay)
        # 添加随机抖动
        jitter = random.uniform(0, delay * 0.1)
        return delay + jitter
    async def connect_with_retry(self):
        """带重试的连接"""
        while self.retry_count < self.max_retries:
            try:
                self.websocket = await websockets.connect(
                    self.uri,
                    ping_interval=20,
                    ping_timeout=10,
                    close_timeout=5
                )
                logger.info(f"Connected successfully after {self.retry_count} retries")
                self.retry_count = 0
                return self.websocket
            except Exception as e:
                self.retry_count += 1
                if self.retry_count >= self.max_retries:
                    logger.error(f"Max retries reached: {e}")
                    raise
                backoff = self.get_backoff_time()
                logger.warning(f"Connection failed (attempt {self.retry_count}): {e}")
                logger.info(f"Retrying in {backoff:.2f} seconds")
                await asyncio.sleep(backoff)
    async def run(self, handler):
        """运行客户端"""
        while True:
            try:
                async with await self.connect_with_retry() as ws:
                    logger.info("WebSocket connected")
                    await handler(ws)
            except websockets.ConnectionClosed:
                logger.warning("Connection closed, reconnecting...")
                self.retry_count = 0
                await asyncio.sleep(1)
            except Exception as e:
                logger.error(f"Unexpected error: {e}")
                await asyncio.sleep(5)
# 使用示例
async def message_handler(websocket):
    """消息处理函数"""
    async for message in websocket:
        logger.info(f"Received: {message}")
        # 处理消息
        await websocket.send(f"Echo: {message}")
async def main():
    client = SmartReconnectWebSocket(
        uri="wss://echo.websocket.org",
        max_retries=10,
        base_delay=1
    )
    await client.run(message_handler)
if __name__ == "__main__":
    asyncio.run(main())
  1. 自动重连:连接断开后自动尝试重新连接
  2. 重连策略:支持固定间隔、指数退避等多种重连策略
  3. 连接状态监控:实时监控连接状态
  4. 消息队列:支持断线期间消息的缓存和重发
  5. 心跳保活:可选的心跳机制保持连接活跃

选择方案时,考虑以下因素:

  • websockets 库:适合 asyncio 异步编程
  • websocket-client 库:适合回调式同步编程
  • 性能要求:高并发场景建议使用异步方案
  • 重连策略:指数退避适合生产环境

抱歉,评论功能暂时关闭!