fix: add a health manager for restarting unhealthy mqtt connections

Author: allenporter

roborock/devices/v1_rpc_channel.py

@@ -13,6 +13,7 @@
 
 from roborock.data import RoborockBase
 from roborock.exceptions import RoborockException
+from roborock.mqtt.health_manager import HealthManager
 from roborock.protocols.v1_protocol import (
     CommandType,
     MapResponse,
@@ -125,12 +126,14 @@ def __init__(
         channel: MqttChannel | LocalChannel,
         payload_encoder: Callable[[RequestMessage], RoborockMessage],
         decoder: Callable[[RoborockMessage], ResponseMessage] | Callable[[RoborockMessage], MapResponse | None],
+        health_manager: HealthManager | None = None,
     ) -> None:
         """Initialize the channel with a raw channel and an encoder function."""
         self._name = name
         self._channel = channel
         self._payload_encoder = payload_encoder
         self._decoder = decoder
+        self._health_manager = health_manager
 
     async def _send_raw_command(
         self,
@@ -165,13 +168,19 @@ def find_response(response_message: RoborockMessage) -> None:
         unsub = await self._channel.subscribe(find_response)
         try:
             await self._channel.publish(message)
-            return await asyncio.wait_for(future, timeout=_TIMEOUT)
+            result = await asyncio.wait_for(future, timeout=_TIMEOUT)
         except TimeoutError as ex:
+            if self._health_manager:
+                await self._health_manager.on_timeout()
             future.cancel()
             raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
         finally:
             unsub()
 
+        if self._health_manager:
+            await self._health_manager.on_success()
+        return result
+
 
 def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:
     """Create a V1 RPC channel using an MQTT channel."""
@@ -180,6 +189,7 @@ def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityDa
         mqtt_channel,
         lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
         decode_rpc_response,
+        health_manager=HealthManager(mqtt_channel.restart),
     )
 
 

roborock/mqtt/health_manager.py

@@ -0,0 +1,51 @@
+"""A health manager for monitoring MQTT connections to Roborock devices.
+
+We observe a problem where sometimes the MQTT connection appears to be alive but
+no messages are being received. To mitigate this, we track consecutive timeouts
+and restart the connection if too many timeouts occur in succession.
+"""
+
+import datetime
+from collections.abc import Awaitable, Callable
+
+# Number of consecutive timeouts before considering the connection unhealthy.
+TIMEOUT_THRESHOLD = 3
+
+# We won't restart the session more often than this interval.
+RESTART_COOLDOWN = datetime.timedelta(minutes=30)
+
+
+class HealthManager:
+    """Manager for monitoring the health of MQTT connections.
+
+    This tracks communication timeouts and can trigger restarts of the MQTT
+    session if too many timeouts occur in succession.
+    """
+
+    def __init__(self, restart: Callable[[], Awaitable[None]]) -> None:
+        """Initialize the health manager.
+
+        Args:
+            restart: A callable to restart the MQTT session.
+        """
+        self._consecutive_timeouts = 0
+        self._restart = restart
+        self._last_restart: datetime.datetime | None = None
+
+    async def on_success(self) -> None:
+        """Record a successful communication event."""
+        self._consecutive_timeouts = 0
+
+    async def on_timeout(self) -> None:
+        """Record a timeout event.
+
+        This may trigger a restart of the MQTT session if too many timeouts
+        have occurred in succession.
+        """
+        self._consecutive_timeouts += 1
+        if self._consecutive_timeouts >= TIMEOUT_THRESHOLD:
+            now = datetime.datetime.now()
+            if self._last_restart is None or now - self._last_restart >= RESTART_COOLDOWN:
+                await self._restart()
+                self._last_restart = now
+                self._consecutive_timeouts = 0

tests/conftest.py

@@ -363,13 +363,17 @@ def __init__(self):
         self.connect = AsyncMock(side_effect=self._connect)
         self.close = MagicMock(side_effect=self._close)
         self.protocol_version = LocalProtocolVersion.V1
+        self.restart = MagicMock(side_effect=self._restart)
 
     async def _connect(self) -> None:
         self._is_connected = True
 
     def _close(self) -> None:
         self._is_connected = False
 
+    def _restart(self) -> None:
+        self._is_connected = False
+
     @property
     def is_connected(self) -> bool:
         """Return true if connected."""

tests/mqtt/test_health_manager.py

@@ -0,0 +1,73 @@
+"""Tests for the health manager."""
+
+import datetime
+from unittest.mock import AsyncMock, patch
+
+from roborock.mqtt.health_manager import HealthManager
+
+
+async def test_health_manager_restart_called_after_timeouts() -> None:
+    """Test that the health manager calls restart after consecutive timeouts."""
+    restart = AsyncMock()
+    health_manager = HealthManager(restart=restart)
+
+    await health_manager.on_timeout()
+    await health_manager.on_timeout()
+    restart.assert_not_called()
+
+    await health_manager.on_timeout()
+    restart.assert_called_once()
+
+
+async def test_health_manager_success_resets_counter() -> None:
+    """Test that a successful message resets the timeout counter."""
+    restart = AsyncMock()
+    health_manager = HealthManager(restart=restart)
+
+    await health_manager.on_timeout()
+    await health_manager.on_timeout()
+    restart.assert_not_called()
+
+    await health_manager.on_success()
+
+    await health_manager.on_timeout()
+    await health_manager.on_timeout()
+    restart.assert_not_called()
+
+    await health_manager.on_timeout()
+    restart.assert_called_once()
+
+
+async def test_cooldown() -> None:
+    """Test that the health manager respects the restart cooldown."""
+    restart = AsyncMock()
+    health_manager = HealthManager(restart=restart)
+
+    with patch("roborock.mqtt.health_manager.datetime") as mock_datetime:
+        now = datetime.datetime(2023, 1, 1, 12, 0, 0)
+        mock_datetime.datetime.now.return_value = now
+
+        # Trigger first restart
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        restart.assert_called_once()
+        restart.reset_mock()
+
+        # Advance time but stay within cooldown (30 mins)
+        mock_datetime.datetime.now.return_value = now + datetime.timedelta(minutes=10)
+
+        # Trigger timeouts again
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        restart.assert_not_called()
+
+        # Advance time past cooldown
+        mock_datetime.datetime.now.return_value = now + datetime.timedelta(minutes=31)
+
+        # Trigger timeouts again
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        await health_manager.on_timeout()
+        restart.assert_called_once()