Shut down MQTT relay and smart plug services explicitly

Existing code did not allow the MQTT relay to be shut down gracefully.

Moved _client.loop_stop() after _client.disconnect() call to allow
Paho to be notified to disconnect and shut down without reconnecting.

Added _client._disconnection_event to help prevent a race condition where
a disconnection is requested but the loop is immediately stopped.
Modify _client.disconnect() to take a timeout parameter (default: 0).
Wait in that routine for the disconnection event with the given timeout.

Apply same logic to MQTTSmartPlugService and add to shutdown sequence.

Shutdown the MQTT relay last in lifespan() with a hard-coded 2 sec timeout.

The disconnect timeout is also available to printer MQTT clients.
It is left to the default of 0 at this time. If later enabled,
it should be a configuration parameters, especially for large farms.

Fixes: #327

Signed-off-by: Jeff Kletsky <git-commits@allycomm.com>

backend/app/main.py

@@ -216,6 +216,7 @@ from backend.app.services.bambu_mqtt import PrinterState
 from backend.app.services.github_backup import github_backup_service
 from backend.app.services.homeassistant import homeassistant_service
 from backend.app.services.mqtt_relay import mqtt_relay
+from backend.app.services.mqtt_smart_plug import mqtt_smart_plug_service
 from backend.app.services.notification_service import notification_service
 from backend.app.services.print_scheduler import scheduler as print_scheduler
 from backend.app.services.printer_manager import (
@@ -2746,6 +2747,10 @@ async def lifespan(app: FastAPI):
     if virtual_printer_manager.is_enabled:
         await virtual_printer_manager.configure(enabled=False)
 
+    await mqtt_smart_plug_service.disconnect(timeout=2)
+
+    await mqtt_relay.disconnect(timeout=2)
+
 
 app = FastAPI(
     title=app_settings.app_name,

backend/app/services/bambu_mqtt.py

@@ -11,6 +11,7 @@ import asyncio
 import json
 import logging
 import ssl
+import threading
 import time
 from collections import deque
 from collections.abc import Callable
@@ -279,6 +280,7 @@ class BambuMQTTClient:
         self._message_log: deque[MQTTLogEntry] = deque(maxlen=100)
         self._logging_enabled: bool = False
         self._last_message_time: float = 0.0  # Track when we last received a message
+        self._disconnection_event: threading.Event | None = None
         self._previous_ams_hash: str | None = None  # Track AMS changes
 
         # K-profile command tracking
@@ -357,6 +359,8 @@ class BambuMQTTClient:
         self.state.connected = False
         if self.on_state_change:
             self.on_state_change(self.state)
+        if self._disconnection_event:
+            self._disconnection_event.set()
 
     def _on_message(self, client, userdata, msg):
         try:
@@ -2371,11 +2375,13 @@ class BambuMQTTClient:
 
         return True
 
-    def disconnect(self):
+    def disconnect(self, timeout: float = 0):
         """Disconnect from the printer."""
         if self._client:
-            self._client.loop_stop()
+            self._disconnection_event = threading.Event()
             self._client.disconnect()
+            self._disconnection_event.wait(timeout=timeout)
+            self._client.loop_stop()
             self._client = None
             self.state.connected = False
 

backend/app/services/mqtt_relay.py

@@ -36,6 +36,7 @@ class MQTTRelayService:
         self._last_printer_status: dict[int, float] = {}  # printer_id -> last publish timestamp
         self._smart_plug_service = None  # Lazy import to avoid circular dependency
         self._settings: dict = {}  # Store settings for smart plug service
+        self._disconnection_event: threading.Event | None = None
 
     async def configure(self, settings: dict) -> bool:
         """Configure MQTT connection from settings.
@@ -187,15 +188,19 @@ class MQTTRelayService:
             logger.warning("MQTT relay disconnected: %s", rc)
         else:
             logger.info("MQTT relay disconnected cleanly")
+        if self._disconnection_event:
+            self._disconnection_event.set()
 
-    async def disconnect(self):
+    async def disconnect(self, timeout: float = 0):
         """Disconnect from MQTT broker."""
         if self.client:
             try:
                 # Publish offline status before disconnecting
                 self._publish_status("offline")
-                self.client.loop_stop()
+                self._disconnection_event = threading.Event()
                 self.client.disconnect()
+                await asyncio.to_thread(self._disconnection_event.wait, timeout=timeout)
+                self.client.loop_stop()
             except Exception as e:
                 logger.debug("MQTT disconnect error (ignored): %s", e)
             finally:

backend/app/services/mqtt_smart_plug.py

@@ -3,6 +3,7 @@
 This service enables integration with Shelly, Zigbee2MQTT, and other MQTT-based energy monitoring devices.
 """
 
+import asyncio
 import json
 import logging
 import threading
@@ -52,6 +53,7 @@ class MQTTSmartPlugService:
         self.plug_configs: dict[int, dict[str, MQTTDataSourceConfig]] = {}
         # plug_id -> latest data
         self.plug_data: dict[int, SmartPlugMQTTData] = {}
+        self._disconnection_event: threading.Event | None = None
         self._configured = False
         self._broker = ""
         self._port = 1883
@@ -209,6 +211,8 @@ class MQTTSmartPlugService:
             logger.warning("MQTT smart plug service disconnected: %s", rc)
         else:
             logger.info("MQTT smart plug service disconnected cleanly")
+        if self._disconnection_event:
+            self._disconnection_event.set()
 
     def _on_message(self, client: mqtt.Client, userdata: Any, msg: mqtt.MQTTMessage):
         """Handle incoming MQTT message, extract data using JSON path."""
@@ -471,12 +475,14 @@ class MQTTSmartPlugService:
         timeout = timedelta(minutes=self.REACHABLE_TIMEOUT_MINUTES)
         return datetime.utcnow() - data.last_seen < timeout
 
-    async def disconnect(self):
+    async def disconnect(self, timeout: float = 0):
         """Disconnect from MQTT broker."""
         if self.client:
             try:
-                self.client.loop_stop()
+                self._disconnection_event = threading.Event()
                 self.client.disconnect()
+                await asyncio.to_thread(self._disconnection_event.wait, timeout=timeout)
+                self.client.loop_stop()
             except Exception as e:
                 logger.debug("MQTT smart plug disconnect error (ignored): %s", e)
             finally:

backend/app/services/printer_manager.py

@@ -209,18 +209,18 @@ class PrinterManager:
         await asyncio.sleep(1)
         return client.state.connected
 
-    def disconnect_printer(self, printer_id: int):
+    def disconnect_printer(self, printer_id: int, timeout: float = 0):
         """Disconnect from a printer."""
         if printer_id in self._clients:
-            self._clients[printer_id].disconnect()
+            self._clients[printer_id].disconnect(timeout=timeout)
             del self._clients[printer_id]
         self._models.pop(printer_id, None)  # Clean up model cache
         self._printer_info.pop(printer_id, None)  # Clean up printer info cache
 
-    def disconnect_all(self):
+    def disconnect_all(self, timeout: float = 0):
         """Disconnect from all printers."""
         for printer_id in list(self._clients.keys()):
-            self.disconnect_printer(printer_id)
+            self.disconnect_printer(printer_id, timeout=timeout)
 
     def get_status(self, printer_id: int) -> PrinterState | None:
         """Get the current status of a printer (checks for stale connections)."""