python实现的websocket日志类

背景

功能需求需要实现一个"实时日志"功能，即在一个可以在web端触发任务的系统中，可以即时显示触发的任务的日志（此后台任务在后台或其他worker主机上执行）。最后采用了websocket方法来实现该功能，即在任务执行端实现一个logger类通过websocket上传实时日志给web后端，再由web后端通过websocket连接传给web前端，实现实时日志的功能。

websocket logHandler类

协程版本

使用websockets库

python 复制代码

import logging
import json
import time
import asyncio
import websockets
from queue import Queue
from threading import Thread
import traceback
from _queue import Empty

class WebSocketHandler(logging.Handler):
    """
    自定义日志处理器，将日志通过WebSocket发送到后端服务
    """
    def __init__(self, ws_url, reconnect_interval=5, max_queue_size=10000):
        """
        初始化WebSocket处理器
        
        Args:
            ws_url: WebSocket服务器URL
            reconnect_interval: 断线重连间隔(秒)
            max_queue_size: 日志队列最大长度，超出时丢弃旧日志
        """
        super().__init__()
        self.ws_url = ws_url
        self.reconnect_interval = reconnect_interval
        self.max_queue_size = max_queue_size
        self.is_running = False
        self.thread = None
        self.queue = None  # 异步队列，在start中初始化
        self.loop = None  # 保存事件循环引用
        
    def emit(self, record):
        """
        重写emit方法，将日志记录发送到WebSocket
        
        Args:
            record: 日志记录对象
        """
        try:
            # 格式化日志
            msg = self.format(record)
            if not msg.endswith("\n"):
                msg += "\n"
            # 跨线程安全添加日志到队列（关键修复）
            if self.loop and self.queue:
                # 使用事件循环的线程安全方法添加元素
                self.loop.call_soon_threadsafe(
                    self._safe_put_queue, msg
                )
            else:
                print("队列未初始化，日志发送失败")
        except Exception as e:
            # 处理发送失败的情况
            self.handleError(record)
    
    def _safe_put_queue(self, msg):
        """线程安全的队列添加方法（在事件循环线程执行）"""
        try:
            if not self.queue.full():
                self.queue.put_nowait(msg)
            else:
                # 队列满时丢弃最旧日志
                self.queue.get_nowait()
                self.queue.put_nowait(msg)
        except Exception as e:
            print(f"队列添加失败: {e}")

    
    def start(self):
        """启动WebSocket发送线程"""
        if not self.is_running:
            self.is_running = True
            self.thread = Thread(target=self._ws_sender_thread)
            self.thread.daemon = True
            self.thread.start()
            print("WebSocket发送线程启动")
    
    def stop(self):
        """停止WebSocket发送线程"""
        self.is_running = False
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2.0)
    
    def _ws_sender_thread(self):
        """WebSocket发送线程主函数"""
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)
        self.queue = asyncio.Queue(maxsize=self.max_queue_size)
        
        async def custom_heartbeat(websocket):
            while True:
                try:
                    await websocket.send(json.dumps({"type":"ping"}))  # 自定义心跳消息
                    await asyncio.sleep(30)  # 30秒间隔
                except Exception as e:
                    print(f"心跳发送失败: {e}")
                    break

        async def _process_logs(websocket):
            # 连接成功后，发送队列中积压的所有日志
            while not self.queue.empty() and self.is_running:
                # log_data = self.queue.get_nowait()
                # await websocket.send(json.dumps(log_data, ensure_ascii=False))
                msg = await self.queue.get()
                await websocket.send(msg)
            
            # 持续发送新日志
            while self.is_running:
                try:
                    # 阻塞等待新日志，带超时以检查线程是否需要停止
                    # log_data = self.queue.get(timeout=1)
                    # await websocket.send(json.dumps(log_data))
                    msg = await self.queue.get()
                    await websocket.send(msg)
                    self.queue.task_done()
                except asyncio.TimeoutError:
                    continue
                except Empty:
                    continue

        async def send_logs():
            while self.is_running:
                try:
                    # 连接WebSocket服务器
                    async with websockets.connect(self.ws_url) as websocket:
                        # 并行运行日志发送和心跳任务
                        await asyncio.gather(
                            _process_logs(websocket),
                            custom_heartbeat(websocket)
                        )
                        
                except Exception as e:
                    traceback.print_exc()
                    # 连接失败或断开，等待后重试
                    # 等待重连间隔（使用异步sleep）
                    await asyncio.sleep(self.reconnect_interval)
        
        try:
            self.loop.run_until_complete(send_logs())
        except Exception as e:
            pass
        finally:
            self.loop.close()
    
    def close(self):
        """关闭处理器"""
        self.stop()
        super().close()

# 配置示例
def setup_logger():
    # 创建logger
    logger = logging.getLogger('my_app')
    logger.setLevel(logging.DEBUG)
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 创建WebSocket处理器
    ws_handler = WebSocketHandler(ws_url='ws://localhost:8999/logs/websocket/client-2546')
    ws_handler.setLevel(logging.DEBUG)
    json_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')  # 我们在handler中自己处理JSON格式化
    ws_handler.setFormatter(json_formatter)
    
    # 添加处理器到logger
    logger.addHandler(console_handler)
    logger.addHandler(ws_handler)
    
    # 启动WebSocket处理器
    ws_handler.start()
    
    return logger

# 使用示例
if __name__ == "__main__":
    logger = setup_logger()
    
    try:
        # 正常记录日志，会同时输出到控制台和WebSocket
        logger.debug('这是一条调试日志')
        logger.info('这是一条信息日志')
        logger.warning('这是一条警告日志')
        logger.error('这是一条错误日志')
        
        # 模拟长时间运行的程序
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        # 程序退出时，确保WebSocket处理器正确关闭
        for handler in logger.handlers:
            if isinstance(handler, WebSocketHandler):
                handler.stop()

多线程版本

使用websocket-client库

python 复制代码

import logging
import json
import time
from queue import Queue, Empty
from threading import Thread
import traceback
import websocket  # 需安装：pip install websocket-client


class WebSocketHandler(logging.Handler):
    """自定义日志处理器（同步版本），通过WebSocket发送日志"""
    def __init__(self, ws_url, reconnect_interval=5, max_queue_size=10000):
        super().__init__()
        self.ws_url = ws_url
        self.reconnect_interval = reconnect_interval
        self.max_queue_size = max_queue_size
        self.is_running = False  # 控制整体运行状态
        self.heartbeat_running = False  # 控制心跳线程
        self.thread = None  # 日志发送主线程
        self.heartbeat_thread = None  # 心跳线程
        self.queue = Queue(maxsize=max_queue_size)  # 同步队列（线程安全）
        self.ws = None  # WebSocket连接实例

    def emit(self, record):
        """日志记录触发时调用，将日志放入队列"""
        try:
            # 格式化日志
            msg = self.format(record)
            if not msg.endswith("\n"):
                msg += "\n"
            
            # 队列满时丢弃最旧日志
            if self.queue.full():
                try:
                    self.queue.get_nowait()  # 移除最旧日志
                except Empty:
                    pass  # 队列已空，无需处理
            self.queue.put_nowait(msg)  # 放入新日志（同步队列线程安全）
        except Exception as e:
            self.handleError(record)
    
    def close(self):
        """关闭处理器"""
        self.stop()
        super().close()

    def start(self):
        """启动日志发送线程和心跳线程"""
        if not self.is_running:
            self.is_running = True
            # 启动日志发送主线程
            self.thread = Thread(target=self._ws_sender_thread)
            self.thread.daemon = True
            self.thread.start()
            print("WebSocket发送线程启动")

    def stop(self):
        """停止所有线程和连接"""
        self.is_running = False
        self.heartbeat_running = False  # 停止心跳线程
        # 关闭WebSocket连接
        if self.ws:
            try:
                self.ws.close()
            except Exception as e:
                print(f"关闭WebSocket失败: {e}")
        # 等待线程结束
        if self.thread and self.thread.is_alive():
            self.thread.join(timeout=2.0)
        if self.heartbeat_thread and self.heartbeat_thread.is_alive():
            self.heartbeat_thread.join(timeout=1.0)
        print("WebSocket发送线程已停止")

    def _start_heartbeat(self):
        """启动心跳线程"""
        self.heartbeat_running = True
        self.heartbeat_thread = Thread(target=self._heartbeat_loop)
        self.heartbeat_thread.daemon = True
        self.heartbeat_thread.start()
        print("心跳线程启动")

    def _heartbeat_loop(self):
        """心跳发送循环（独立线程）"""
        while self.heartbeat_running and self.is_running:
            try:
                if self.ws and self.ws.connected:  # 检查连接是否有效
                    self.ws.send(json.dumps({"type": "ping"}))  # 发送心跳
                time.sleep(30)  # 30秒间隔
            except Exception as e:
                print(f"心跳发送失败: {e}")
                break  # 心跳失败，退出循环（由主线程重连）

    def _process_logs(self):
        """处理队列中的日志并发送（同步阻塞）"""
        while self.is_running:
            try:
                # 阻塞等待日志（超时1秒，避免永久阻塞）
                msg = self.queue.get(timeout=1)
                if self.ws and self.ws.connected:
                    self.ws.send(msg)  # 发送日志
                    self.queue.task_done()
                else:
                    # 连接已断开，将日志放回队列
                    self.queue.put(msg)
                    time.sleep(0.1)  # 短暂等待后重试
            except Empty:
                continue  # 队列空，继续循环
            except Exception as e:
                print(f"日志发送失败: {e}")
                # 发送失败，将日志放回队列重试
                try:
                    self.queue.put(msg)
                except Exception as put_err:
                    print(f"日志放回队列失败: {put_err}")
                time.sleep(1)  # 等待后重试

    def _ws_sender_thread(self):
        """WebSocket发送主线程：负责连接和日志发送协调"""
        while self.is_running:
            try:
                # 建立WebSocket连接
                print(f"连接WebSocket服务器: {self.ws_url}")
                self.ws = websocket.create_connection(self.ws_url)
                print("WebSocket连接成功")

                # 启动心跳线程（每次重连后重启心跳）
                self._start_heartbeat()

                # 处理日志发送
                self._process_logs()

            except Exception as e:
                print(f"WebSocket连接/发送异常: {e}")
                traceback.print_exc()

            finally:
                # 连接断开时清理
                self.heartbeat_running = False  # 停止当前心跳线程
                if self.heartbeat_thread:
                    self.heartbeat_thread.join(timeout=1.0)
                if self.ws:
                    try:
                        self.ws.close()
                    except Exception as e:
                        print(f"关闭WebSocket连接失败: {e}")
                self.ws = None  # 重置连接实例

                # 断线重连等待
                if self.is_running:
                    print(f"等待{self.reconnect_interval}秒后重试...")
                    time.sleep(self.reconnect_interval)


# 配置示例
def setup_logger():
    # 创建logger
    logger = logging.getLogger('my_app')
    logger.setLevel(logging.DEBUG)
    
    # 创建控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    console_handler.setFormatter(console_formatter)
    
    # 创建WebSocket处理器（使用同步版本）
    ws_handler = WebSocketHandler(ws_url='ws://localhost:8999/logs/websocket/client-2546')
    ws_handler.setLevel(logging.DEBUG)
    ws_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    ws_handler.setFormatter(ws_formatter)
    
    # 添加处理器到logger
    logger.addHandler(console_handler)
    logger.addHandler(ws_handler)
    
    # 启动WebSocket处理器
    ws_handler.start()
    
    return logger


# 使用示例
if __name__ == "__main__":
    logger = setup_logger()
    
    try:
        # 测试日志发送
        logger.debug('这是一条调试日志')
        logger.info('这是一条信息日志')
        logger.warning('这是一条警告日志')
        logger.error('这是一条错误日志')
        
        # 模拟长时间运行的程序
        while True:
            logger.info('持续发送的日志...')
            time.sleep(5)  # 每5秒发送一条测试日志
    except KeyboardInterrupt:
        print("程序退出中...")
        # 停止WebSocket处理器
        for handler in logger.handlers:
            if isinstance(handler, WebSocketHandler):
                handler.stop()
        print("程序已退出")

集成

集成时只需将handler的示例加到全局logger中即可，就像main函数中setup_logger()的使用那样，但需注意正式使用时最好将handler.stop()函数放在finally块中，确保正确退出。