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) |