dji_simulator/src/core/connection_manager.py

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