118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
import paho.mqtt.client as mqtt
|
|
import logging
|
|
import time
|
|
from threading import Thread, Event
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.client = None
|
|
self.connected = False
|
|
self._stop_event = Event()
|
|
self._heartbeat_thread = None
|
|
self._message_handler = None
|
|
|
|
def set_message_handler(self, handler):
|
|
self._message_handler = handler
|
|
|
|
def connect(self, broker_ip: str, client_id: str, port: int = 1883, username: str = None, password: str = None):
|
|
"""
|
|
Connects to the MQTT Broker.
|
|
"""
|
|
if self.client:
|
|
self.disconnect()
|
|
|
|
# Create a new MQTT Client instance
|
|
# Protocol V5 is recommended for modern deployments, but V3.1.1 is standard.
|
|
# Using default (usually 3.1.1 or 5.0 depending on paho version/negotiation).
|
|
self.client = mqtt.Client(client_id=client_id, protocol=mqtt.MQTTv311)
|
|
|
|
if username and password:
|
|
self.client.username_pw_set(username, password)
|
|
|
|
# Set callbacks
|
|
self.client.on_connect = self._on_connect
|
|
self.client.on_disconnect = self._on_disconnect
|
|
self.client.on_message = self._on_message
|
|
self.client.on_publish = self._on_publish
|
|
|
|
logger.info(f"Connecting to {broker_ip}:{port} as {client_id}...")
|
|
try:
|
|
self.client.connect(broker_ip, port, keepalive=60)
|
|
self.client.loop_start() # Start the network loop in a background thread
|
|
except Exception as e:
|
|
logger.error(f"Failed to connect: {e}")
|
|
raise
|
|
|
|
def disconnect(self):
|
|
"""
|
|
Disconnects from the broker.
|
|
"""
|
|
self._stop_heartbeat()
|
|
if self.client:
|
|
self.client.loop_stop()
|
|
self.client.disconnect()
|
|
self.client = None
|
|
self.connected = False
|
|
logger.info("Disconnected.")
|
|
|
|
def _on_connect(self, client, userdata, flags, rc):
|
|
if rc == 0:
|
|
self.connected = True
|
|
logger.info("Connected successfully.")
|
|
self._start_heartbeat()
|
|
# Subscribe to services if needed, but usually caller handles subscription
|
|
# or we do it here if we know the topic.
|
|
# For now, we rely on caller or hardcoded.
|
|
# Let's subscribe to everything for this simulator or specific service topic
|
|
# client.subscribe("#")
|
|
# Ideally, we should let the ServiceHandler subscribe.
|
|
else:
|
|
self.connected = False
|
|
logger.error(f"Connection failed with result code {rc}")
|
|
|
|
def _on_disconnect(self, client, userdata, rc):
|
|
self.connected = False
|
|
self._stop_heartbeat()
|
|
if rc != 0:
|
|
logger.warning("Unexpected disconnection.")
|
|
else:
|
|
logger.info("Disconnected gracefully.")
|
|
|
|
def _on_message(self, client, userdata, msg):
|
|
logger.debug(f"Received message on {msg.topic}: {msg.payload}")
|
|
if self._message_handler:
|
|
self._message_handler(msg.topic, msg.payload)
|
|
|
|
def _on_publish(self, client, userdata, mid):
|
|
pass
|
|
|
|
def _start_heartbeat(self):
|
|
self._stop_event.clear()
|
|
self._heartbeat_thread = Thread(target=self._heartbeat_loop, daemon=True)
|
|
self._heartbeat_thread.start()
|
|
|
|
def _stop_heartbeat(self):
|
|
self._stop_event.set()
|
|
if self._heartbeat_thread:
|
|
self._heartbeat_thread.join(timeout=1.0)
|
|
self._heartbeat_thread = None
|
|
|
|
def _heartbeat_loop(self):
|
|
logger.info("Heartbeat loop started.")
|
|
while not self._stop_event.is_set():
|
|
if self.client and self.connected:
|
|
# Actual topic/payload depends on DJI API specs.
|
|
# For now, we simulate a ping or minimal publish.
|
|
# Usually MQTT keepalive handles PINGREQ/PINGRESP automatically.
|
|
# But sometimes application-level heartbeat is needed.
|
|
# We will rely on Paho's built-in keepalive for network ping,
|
|
# but if an app-level heartbeat is needed (e.g. reporting state), we do it here.
|
|
# For now, let's just log or sleep.
|
|
pass
|
|
time.sleep(5)
|
|
|
|
def is_connected(self):
|
|
return self.connected
|