248 lines
9.0 KiB
Python
248 lines
9.0 KiB
Python
|
#!/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)
|