256 lines
11 KiB
Python
Raw Normal View History

2025-09-29 09:35:08 +08:00
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
MQTT相关功能模块
"""
import queue
import time
from typing import Optional, Dict, Any
from utils.logger import get_logger
import paho.mqtt.client as mqtt
logger = get_logger("services.online_script.built_in_modules.mqtt_module")
class VWEDMqttModule:
"""MQTT模块 - 提供MQTT订阅和发布功能"""
def __init__(self, script_id: str):
self.script_id = script_id
self._clients: Dict[str, mqtt.Client] = {}
self._message_queues: Dict[str, queue.Queue] = {}
self._config = self._load_config()
self._connected_clients: Dict[str, bool] = {}
def _load_config(self) -> Dict[str, Any]:
"""加载MQTT配置"""
return {
'enable': True,
'pub_config': { # 发布者
'broker': 'tcp://192.168.189.97:1883', # 服务器端点url
'topics': ['oagv/v2/asbm2_IRAYPLE/SIM-1/state'], # 订阅的主题MQTT允许使用通配符订阅主题不允许使用通配符pub发布消息
'qos': 2, # 设置message的服务质量0消息最多传递一次零次或一次1至少传递一次一次或多次。2只传递一次
'client_id': 'RDS-Pub', # 客户端唯一标识
'username': None, # 连接的用户名,不需要就填 null
'password': None, # 连接的密码,不需要就填 null
'clean_session': False, # 设置是否清空session,false表示服务器会保留客户端的连接记录true每次都以新的身份连接服务器
'connection_timeout': 30, # 超时时间(seconds)
'keep_alive_interval': 60, # 设置会话心跳时间(seconds)
'automatic_reconnect': True, # 设置断开后重新连接
'retained': False, # 表示发送的消息需要一直持久保存不受服务器重启影响不但要发送给当前的订阅者并且以后新来的订阅了此Topic name的订阅者会马上得到推送。
'will_msg': None, # 作为publish的遗嘱消息会存到服务器在publish端非正常断连的情况下发送给所有订阅的客户端 不需要就填 null
'will_topic': 'Examples1' # 遗嘱消息发布的topic
},
'sub_config': { # 订阅者
'broker': 'tcp://192.168.189.97:1883', # 服务器端点url
'topics': ['uagv/v2/IRAYPLE/SIM-1/state'], # MQTT允许使用通配符sub订阅主题但不允许使用通配符pub发布消息
'qos': 1, # 设置message的服务质量0消息最多传递一次零次或一次1至少传递一次一次或多次。2只传递一次
'client_id': 'RDS-Sub', # 客户端唯一标识
'username': None, # 连接的用户名,不需要就填 null
'password': None, # 连接的密码,不需要就填 null
'clean_session': False, # 设置是否清空session,false表示服务器会保留客户端的连接记录true每次都以新的身份连接服务器
'connection_timeout': 30, # 超时时间(seconds)
'keep_alive_interval': 60, # 设置会话心跳时间(seconds)
'automatic_reconnect': True, # 设置断开后重新连接
'retained': False # 表示发送的消息需要一直持久保存不受服务器重启影响不但要发送给当前的订阅者并且以后新来的订阅了此Topic name的订阅者会马上得到推送。
}
}
def _parse_broker_url(self, broker: str) -> tuple:
"""解析broker URL"""
if broker.startswith('tcp://'):
broker = broker[6:]
elif broker.startswith('mqtt://'):
broker = broker[7:]
if ':' in broker:
host, port = broker.split(':', 1)
return host, int(port)
else:
return broker, 1883
def _get_or_create_client(self, client_type: str, config: Dict[str, Any]) -> Optional[mqtt.Client]:
"""获取或创建MQTT客户端"""
client_key = f"{client_type}_{self.script_id}"
if client_key in self._clients:
return self._clients[client_key]
try:
client = mqtt.Client(
client_id=config['client_id'] + f"_{self.script_id}",
callback_api_version=mqtt.CallbackAPIVersion.VERSION2
)
if config.get('username') and config.get('password'):
client.username_pw_set(config['username'], config['password'])
client.clean_session = config.get('clean_session', False)
if client_type == 'sub':
if client_key not in self._message_queues:
self._message_queues[client_key] = queue.Queue()
def on_message(_client, _userdata, msg, _properties=None):
try:
message = msg.payload.decode('utf-8')
self._message_queues[client_key].put(message)
except Exception as e:
logger.error(f"Error processing MQTT message: {e}")
client.on_message = on_message
def on_connect(_client, _userdata, _flags, rc, _properties=None):
if rc == 0:
self._connected_clients[client_key] = True
logger.info(f"MQTT client {client_key} connected successfully")
else:
self._connected_clients[client_key] = False
logger.error(f"MQTT client {client_key} failed to connect with code {rc}")
def on_disconnect(_client, _userdata, rc, _properties=None):
self._connected_clients[client_key] = False
logger.info(f"MQTT client {client_key} disconnected")
client.on_connect = on_connect
client.on_disconnect = on_disconnect
host, port = self._parse_broker_url(config['broker'])
client.connect(host, port, config.get('keep_alive_interval', 60))
client.loop_start()
self._clients[client_key] = client
# 等待连接建立
timeout = time.time() + config.get('connection_timeout', 30)
while time.time() < timeout and not self._connected_clients.get(client_key, False):
time.sleep(0.1)
if not self._connected_clients.get(client_key, False):
logger.error(f"MQTT client {client_key} connection timeout")
return None
return client
except Exception as e:
logger.error(f"Failed to create MQTT client: {e}")
return None
def mqtt_subscribe(self, topic: Optional[str] = None) -> Optional[str]:
"""MQTT订阅信息
Args:
topic: 订阅的topic如果不提供则使用配置文件中的topic
Returns:
订阅到的信息失败返回None
"""
if not self._config.get('enable', False):
logger.warning("MQTT is disabled in configuration")
return None
sub_config = self._config.get('sub_config', {})
client = self._get_or_create_client('sub', sub_config)
if not client:
return None
try:
# 确定要订阅的topic
if topic:
subscribe_topic = topic
else:
topics = sub_config.get('topics', [])
if not topics:
logger.error("No topics configured for subscription")
return None
subscribe_topic = topics[0] # 使用第一个配置的topic
print("subscribe_topic::", subscribe_topic)
# 订阅topic
client.subscribe(subscribe_topic, sub_config.get('qos', 1))
# 等待消息
client_key = f"sub_{self.script_id}"
message_queue = self._message_queues.get(client_key)
if not message_queue:
logger.error("Message queue not available")
return None
try:
# 等待消息超时时间5秒
message = message_queue.get(timeout=5)
return message
except queue.Empty:
logger.warning("No message received within timeout")
return None
except Exception as e:
logger.error(f"MQTT subscribe failed: {e}")
return None
def mqtt_publish(self, topic_or_message: str, message: Optional[str] = None) -> bool:
"""MQTT发布信息
Args:
topic_or_message: 如果message为None则此参数为消息内容否则为topic
message: 消息内容
Returns:
发布是否成功
"""
if not self._config.get('enable', False):
logger.warning("MQTT is disabled in configuration")
return False
pub_config = self._config.get('pub_config', {})
client = self._get_or_create_client('pub', pub_config)
if not client:
return False
try:
# 确定topic和message
if message is None:
# 使用配置文件中的topic
topics = pub_config.get('topics', [])
if not topics:
logger.error("No topics configured for publishing")
return False
publish_topic = topics[0]
publish_message = topic_or_message
else:
# 使用传入的topic
publish_topic = topic_or_message
publish_message = message
# 发布消息
result = client.publish(
publish_topic,
publish_message,
qos=pub_config.get('qos', 2),
retain=pub_config.get('retained', False)
)
return result.rc == mqtt.MQTT_ERR_SUCCESS
except Exception as e:
logger.error(f"MQTT publish failed: {e}")
return False
def cleanup(self):
"""清理资源"""
for client in self._clients.values():
try:
client.loop_stop()
client.disconnect()
except:
pass
self._clients.clear()
self._message_queues.clear()
self._connected_clients.clear()