fix: Share a HealthManager instance across all mqtt channels (#672)

The motiviation is to better isolate single device timeouts from overall session health. I have not seen this as an explicit issue, but seems worth being defensive against. This will ensure that a single inactive device will not unnecessarily restart the MQTT session. Any successful RPCs by other devices will reset the counter.

Adds explicit logging when a session is reset by the health manager.

Author: allenporter

roborock/devices/mqtt_channel.py

@@ -6,6 +6,7 @@
 from roborock.callbacks import decoder_callback
 from roborock.data import HomeDataDevice, RRiot, UserData
 from roborock.exceptions import RoborockException
+from roborock.mqtt.health_manager import HealthManager
 from roborock.mqtt.session import MqttParams, MqttSession, MqttSessionException
 from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder
 from roborock.roborock_message import RoborockMessage
@@ -42,6 +43,11 @@ def is_connected(self) -> bool:
         """
         return self._mqtt_session.connected
 
+    @property
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._mqtt_session.health_manager
+
     @property
     def is_local_connected(self) -> bool:
         """Return true if the channel is connected locally."""

roborock/devices/v1_channel.py

@@ -181,7 +181,6 @@ def __init__(
         self._logger = RoborockLoggerAdapter(device_uid, _LOGGER)
         self._security_data = security_data
         self._mqtt_channel = mqtt_channel
-        self._mqtt_health_manager = HealthManager(self._mqtt_channel.restart)
         self._local_session = local_session
         self._local_channel: LocalChannel | None = None
         self._mqtt_unsub: Callable[[], None] | None = None
@@ -272,7 +271,7 @@ def _create_mqtt_rpc_strategy(self, decoder: Callable[[RoborockMessage], Any] =
                 security_data=self._security_data,
             ),
             decoder=decoder,
-            health_manager=self._mqtt_health_manager,
+            health_manager=self._mqtt_channel.health_manager,
         )
 
     async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:

roborock/mqtt/health_manager.py

@@ -6,8 +6,11 @@
 """
 
 import datetime
+import logging
 from collections.abc import Awaitable, Callable
 
+_LOGGER = logging.getLogger(__name__)
+
 # Number of consecutive timeouts before considering the connection unhealthy.
 TIMEOUT_THRESHOLD = 3
 
@@ -45,7 +48,13 @@ async def on_timeout(self) -> None:
         self._consecutive_timeouts += 1
         if self._consecutive_timeouts >= TIMEOUT_THRESHOLD:
             now = datetime.datetime.now(datetime.UTC)
-            if self._last_restart is None or now - self._last_restart >= RESTART_COOLDOWN:
+            since_last = (now - self._last_restart) if self._last_restart else None
+            if since_last is None or since_last >= RESTART_COOLDOWN:
+                _LOGGER.debug(
+                    "Restarting MQTT session after %d consecutive timeouts (duration since last restart %s)",
+                    self._consecutive_timeouts,
+                    since_last,
+                )
                 await self._restart()
                 self._last_restart = now
                 self._consecutive_timeouts = 0

roborock/mqtt/roborock_session.py

@@ -19,6 +19,7 @@
 
 from roborock.callbacks import CallbackMap
 
+from .health_manager import HealthManager
 from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized
 
 _LOGGER = logging.getLogger(__name__)
@@ -75,12 +76,18 @@ def __init__(
         self._connection_task: asyncio.Task[None] | None = None
         self._topic_idle_timeout = topic_idle_timeout
         self._idle_timers: dict[str, asyncio.Task[None]] = {}
+        self._health_manager = HealthManager(self.restart)
 
     @property
     def connected(self) -> bool:
         """True if the session is connected to the broker."""
         return self._healthy
 
+    @property
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._health_manager
+
     async def start(self) -> None:
         """Start the MQTT session.
 
@@ -337,6 +344,11 @@ def connected(self) -> bool:
         """True if the session is connected to the broker."""
         return self._session.connected
 
+    @property
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._session.health_manager
+
     async def _maybe_start(self) -> None:
         """Start the MQTT session if not already started."""
         async with self._lock:

roborock/mqtt/session.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 
 from roborock.exceptions import RoborockException
+from roborock.mqtt.health_manager import HealthManager
 
 DEFAULT_TIMEOUT = 30.0
 
@@ -40,6 +41,11 @@ class MqttSession(ABC):
     def connected(self) -> bool:
         """True if the session is connected to the broker."""
 
+    @property
+    @abstractmethod
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+
     @abstractmethod
     async def subscribe(self, device_id: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
         """Invoke the callback when messages are received on the topic.

tests/conftest.py

@@ -13,6 +13,7 @@
 
 from roborock import HomeData, UserData
 from roborock.data import DeviceData
+from roborock.mqtt.health_manager import HealthManager
 from roborock.protocols.v1_protocol import LocalProtocolVersion
 from roborock.roborock_message import RoborockMessage
 from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1
@@ -364,6 +365,7 @@ def __init__(self):
         self.close = MagicMock(side_effect=self._close)
         self.protocol_version = LocalProtocolVersion.V1
         self.restart = AsyncMock()
+        self.health_manager = HealthManager(self.restart)
 
     async def _connect(self) -> None:
         self._is_connected = True