#!/usr/bin/env python # -*- coding: utf-8 -*- """ 异步MQTT客户端服务模块 提供高性能的异步MQTT连接、发布、订阅等功能 支持自动重连和异步消息处理 """ import asyncio import json import time import sys from typing import Dict, Any, Optional, Callable, List from utils.logger import get_logger # Windows兼容性修复 if sys.platform == 'win32': # 在Windows上使用SelectorEventLoop以支持add_reader/add_writer asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) # try: # import aiomqtt # except ImportError: # try: # # 尝试导入旧版本的asyncio-mqtt # # except ImportError: # aiomqtt = None # 如果都没有安装，回退到paho-mqtt import aiomqtt from config.tf_api_config import MQTT_CONFIG logger = get_logger("services.async_mqtt_service") class AsyncMQTTService: """异步MQTT客户端服务""" def __init__(self, config: Dict[str, Any] = None): self.config = config or MQTT_CONFIG self.client = None self.connected = False self.reconnect_attempts = 0 self.message_handlers: Dict[str, List[Callable]] = {} self._running = False self._connection_task = None self._message_task = None self._use_aiomqtt = aiomqtt is not None if not self._use_aiomqtt: logger.warning("aiomqtt未安装，将使用paho-mqtt的异步包装") async def connect(self): """连接MQTT服务器""" try: if self._use_aiomqtt: await self._connect_aiomqtt() else: await self._connect_paho_async() return True except Exception as e: logger.error(f"连接MQTT服务器失败: {e}") return False async def _connect_aiomqtt(self): """使用aiomqtt连接""" client_config = { "hostname": self.config['host'], "port": self.config['port'], # "keepalive": self.config.get('keepalive', 60), } if self.config.get('username'): client_config['username'] = self.config['username'] client_config['password'] = self.config.get('password', '') # try: self.client = aiomqtt.Client(**client_config) # except TypeError as e: # # 处理版本兼容性问题，移除不支持的参数 # logger.warning(f"aiomqtt版本兼容性问题，重试连接: {e}") # # 创建基础配置，只包含核心参数 # basic_config = { # "hostname": self.config['host'], # "port": self.config['port'] # } # if self.config.get('username'): # basic_config['username'] = self.config['username'] # basic_config['password'] = self.config.get('password', '') # self.client = aiomqtt.Client(**basic_config) self._running = True self._connection_task = asyncio.create_task(self._maintain_connection()) logger.info(f"正在连接MQTT服务器 {self.config['host']}:{self.config['port']}...") # 等待连接建立 timeout = 10 start_time = time.time() while not self.connected and (time.time() - start_time) < timeout: await asyncio.sleep(0.1) if self.connected: logger.info("异步MQTT连接建立成功") else: logger.error("异步MQTT连接建立超时") async def _maintain_connection(self): """维护MQTT连接（仅用于aiomqtt）""" while self._running: try: async with self.client: self.connected = True self.reconnect_attempts = 0 logger.info(f"MQTT连接成功: {self.config['host']}:{self.config['port']}") # 启动消息处理任务 self._message_task = asyncio.create_task(self._handle_messages()) # 重新订阅所有topics await self._resubscribe_all_topics() # 等待消息处理任务完成或连接断开 await self._message_task except Exception as e: self.connected = False if self._message_task: self._message_task.cancel() if self._running: self.reconnect_attempts += 1 if self.reconnect_attempts <= self.config.get('max_retries', 3): delay = self.config.get('reconnect_delay', 5) * self.reconnect_attempts logger.warning(f"MQTT连接断开: {e}, 将在 {delay} 秒后重连（第{self.reconnect_attempts}次尝试）") await asyncio.sleep(delay) else: logger.error("MQTT重连次数已达上限，停止重连") break else: break async def _handle_messages(self): """处理接收到的消息（仅用于aiomqtt）""" try: async for message in self.client.messages: topic = str(message.topic) payload = message.payload.decode('utf-8') # 异步处理消息 await self._process_message(topic, payload) except asyncio.CancelledError: logger.debug("消息处理任务被取消") except Exception as e: logger.error(f"处理MQTT消息失败: {e}", exc_info=True) async def _process_message(self, topic: str, payload: str): """异步处理单个消息""" try: handlers = self.message_handlers.get(topic, []) for handler in handlers: try: if asyncio.iscoroutinefunction(handler): await handler(topic, payload) else: # 对于同步处理器，在线程池中执行 loop = asyncio.get_event_loop() await loop.run_in_executor(None, handler, topic, payload) except Exception as e: logger.error(f"MQTT消息处理器执行失败: {e}", exc_info=True) except Exception as e: logger.error(f"处理MQTT消息异常: {e}", exc_info=True) async def _resubscribe_all_topics(self): """重新订阅所有topics""" for topic in self.message_handlers.keys(): try: await self.client.subscribe(topic) logger.debug(f"重新订阅MQTT主题: {topic}") except Exception as e: logger.error(f"重新订阅主题失败 {topic}: {e}") async def disconnect(self): """断开连接""" self._running = False self.connected = False if self._message_task: self._message_task.cancel() try: await self._message_task except asyncio.CancelledError: pass if self._connection_task: self._connection_task.cancel() try: await self._connection_task except asyncio.CancelledError: pass logger.info("异步MQTT连接已断开") async def subscribe(self, topic: str, handler: Callable = None): """订阅主题""" try: if not self.connected: logger.warning(f"MQTT未连接，无法订阅主题: {topic}") return False if self._use_aiomqtt and self.client: await self.client.subscribe(topic) logger.info(f"成功订阅MQTT主题: {topic}") # 添加消息处理器 if handler: if topic not in self.message_handlers: self.message_handlers[topic] = [] if handler not in self.message_handlers[topic]: self.message_handlers[topic].append(handler) return True else: logger.error("MQTT客户端未初始化") return False except Exception as e: logger.error(f"订阅MQTT主题异常: {e}") return False async def unsubscribe(self, topic: str): """取消订阅主题""" try: if not self.connected: logger.debug(f"MQTT未连接，跳过取消订阅主题: {topic}") return False if self._use_aiomqtt and self.client: await self.client.unsubscribe(topic) logger.info(f"取消订阅MQTT主题: {topic}") # 清理消息处理器 if topic in self.message_handlers: del self.message_handlers[topic] return True else: logger.warning(f"MQTT客户端未初始化，无法取消订阅主题: {topic}") return False except Exception as e: logger.error(f"取消订阅MQTT主题异常: {e}") return False async def publish(self, topic: str, payload: Any, qos: int = 0, retain: bool = False): """发布消息""" try: if not self.connected: logger.warning(f"MQTT未连接，无法发布消息到: {topic}") return False # 处理不同类型的payload payload_str = self._serialize_payload(payload) if self._use_aiomqtt and self.client: await self.client.publish(topic, payload_str, qos=qos, retain=retain) logger.debug(f"MQTT消息发布成功: {topic}") return True else: logger.error("MQTT客户端未初始化") return False except Exception as e: logger.error(f"发布MQTT消息异常: {e}") return False def _serialize_payload(self, payload: Any) -> str: """序列化消息载荷""" try: if isinstance(payload, (str, int, float, bool)): return str(payload) elif isinstance(payload, dict): return json.dumps(payload, ensure_ascii=False, default=self._json_default) elif isinstance(payload, list): return json.dumps(payload, ensure_ascii=False, default=self._json_default) elif hasattr(payload, 'to_dict') and callable(payload.to_dict): # 如果对象有to_dict方法，优先使用它 return json.dumps(payload.to_dict(), ensure_ascii=False, default=self._json_default) elif hasattr(payload, 'dict'): # 如果是Pydantic模型等，尝试使用dict方法 if callable(payload.dict): return json.dumps(payload.dict(), ensure_ascii=False, default=self._json_default) else: return json.dumps(payload.dict, ensure_ascii=False, default=self._json_default) elif hasattr(payload, '__dict__'): # 对于自定义对象，尝试序列化其属性 return json.dumps(payload.__dict__, ensure_ascii=False, default=self._json_default) else: # 最后尝试直接转换为字符串 return str(payload) except Exception as e: logger.warning(f"序列化载荷失败，使用字符串表示: {e}") return str(payload) def _json_default(self, obj): """JSON序列化的默认处理器，处理特殊类型""" from enum import Enum if isinstance(obj, Enum): return obj.value elif hasattr(obj, 'to_dict') and callable(obj.to_dict): return obj.to_dict() elif hasattr(obj, '__dict__'): return obj.__dict__ else: return str(obj) def add_message_handler(self, topic: str, handler: Callable): """为指定主题添加消息处理器""" if topic not in self.message_handlers: self.message_handlers[topic] = [] if handler not in self.message_handlers[topic]: self.message_handlers[topic].append(handler) def remove_message_handler(self, topic: str, handler: Callable): """移除指定主题的消息处理器""" if topic in self.message_handlers: if handler in self.message_handlers[topic]: self.message_handlers[topic].remove(handler) if not self.message_handlers[topic]: del self.message_handlers[topic] def is_connected(self) -> bool: """检查连接状态""" return self.connected def get_connection_info(self) -> Dict[str, Any]: """获取连接信息""" return { "host": self.config['host'], "port": self.config['port'], "connected": self.connected, "reconnect_attempts": self.reconnect_attempts, "running": self._running, "subscribed_topics": list(self.message_handlers.keys()), "async_mode": self._use_aiomqtt } def create_async_mqtt_service(config: Dict[str, Any] = None) -> AsyncMQTTService: """创建异步MQTT服务实例""" return AsyncMQTTService(config)