256 lines
11 KiB
Python
256 lines
11 KiB
Python
#!/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() |