#!/usr/bin/env python # -*- coding: utf-8 -*- """ MQTT客户端服务模块 提供MQTT连接、发布、订阅等功能 支持自动重连和消息处理 """ import asyncio import json import time import threading from typing import Dict, Any, Optional, Callable, List from utils.logger import get_logger import paho.mqtt.client as mqtt from config.tf_api_config import MQTT_CONFIG logger = get_logger("services.mqtt_service") class MQTTService: """内部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._lock = threading.RLock() self._running = False def on_connect(self, client, userdata, flags, rc): """连接回调""" if rc == 0: self.connected = True self.reconnect_attempts = 0 logger.info(f"MQTT连接成功: {self.config['host']}:{self.config['port']}") else: self.connected = False logger.error(f"MQTT连接失败，错误代码: {rc}") def on_disconnect(self, client, userdata, rc): """断开连接回调""" self.connected = False logger.warning(f"MQTT连接断开，代码: {rc}") if self._running and rc != 0: self._schedule_reconnect() def on_message(self, client, userdata, msg): """消息接收回调""" try: topic = msg.topic payload = msg.payload.decode('utf-8') with self._lock: handlers = self.message_handlers.get(topic, []) for handler in handlers: try: if asyncio.iscoroutinefunction(handler): asyncio.create_task(handler(topic, payload)) else: 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) def _schedule_reconnect(self): """安排重连""" if self.reconnect_attempts >= self.config.get('max_retries', 3): logger.error("MQTT重连次数已达上限，停止重连") return delay = self.config.get('reconnect_delay', 5) * (self.reconnect_attempts + 1) self.reconnect_attempts += 1 logger.info(f"将在 {delay} 秒后重连MQTT（第{self.reconnect_attempts}次尝试）") def reconnect(): time.sleep(delay) if self._running: self.connect() threading.Thread(target=reconnect, daemon=True).start() def connect(self): """连接MQTT服务器""" try: if self.client: self.client.disconnect() self.client = mqtt.Client() self.client.on_connect = self.on_connect self.client.on_message = self.on_message self.client.on_disconnect = self.on_disconnect # 设置认证信息 if self.config.get('username'): self.client.username_pw_set( self.config['username'], self.config.get('password', '') ) logger.info(f"正在连接MQTT服务器 {self.config['host']}:{self.config['port']}...") self.client.connect( self.config['host'], self.config['port'], self.config.get('keepalive', 60) ) self.client.loop_start() self._running = True # 等待连接建立 timeout = 10 start_time = time.time() while not self.connected and (time.time() - start_time) < timeout: time.sleep(0.1) if self.connected: logger.info("MQTT连接建立成功") return True else: logger.error("MQTT连接建立超时") return False except Exception as e: logger.error(f"连接MQTT服务器失败: {e}") return False def disconnect(self): """断开连接""" self._running = False if self.client: self.client.loop_stop() self.client.disconnect() self.connected = False logger.info("MQTT连接已断开") def subscribe(self, topic: str, handler: Callable = None): """订阅主题""" try: if not self.client or not self.connected: logger.warning(f"MQTT未连接，无法订阅主题: {topic}") return False result = self.client.subscribe(topic) if result[0] == mqtt.MQTT_ERR_SUCCESS: logger.info(f"成功订阅MQTT主题: {topic}") # 添加消息处理器 if handler: with self._lock: 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(f"订阅MQTT主题失败: {topic}, 错误码: {result[0]}") return False except Exception as e: logger.error(f"订阅MQTT主题异常: {e}") return False def unsubscribe(self, topic: str): """取消订阅主题""" try: if self.client and self.connected: result = self.client.unsubscribe(topic) if result[0] == mqtt.MQTT_ERR_SUCCESS: logger.info(f"取消订阅MQTT主题: {topic}") # 清理消息处理器 with self._lock: if topic in self.message_handlers: del self.message_handlers[topic] return True else: logger.error(f"取消订阅MQTT主题失败: {topic}, 错误码: {result[0]}") return False except Exception as e: logger.error(f"取消订阅MQTT主题异常: {e}") return False def publish(self, topic: str, payload: Any, qos: int = 0, retain: bool = False): """发布消息""" try: if not self.client or not self.connected: logger.warning(f"MQTT未连接，无法发布消息到: {topic}") return False if isinstance(payload, dict): payload_str = json.dumps(payload, ensure_ascii=False) else: payload_str = str(payload) result = self.client.publish(topic, payload_str, qos, retain) if result.rc == mqtt.MQTT_ERR_SUCCESS: logger.debug(f"MQTT消息发布成功: {topic}") return True else: logger.error(f"MQTT消息发布失败: {topic}, 错误码: {result.rc}") return False except Exception as e: logger.error(f"发布MQTT消息异常: {e}") return False def add_message_handler(self, topic: str, handler: Callable): """为指定主题添加消息处理器""" with self._lock: 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): """移除指定主题的消息处理器""" with self._lock: 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()) } def create_mqtt_service(config: Dict[str, Any] = None) -> MQTTService: """创建MQTT服务实例""" return MQTTService(config)