fix: adding keepalive to local connection

Author: humbertogontijo

roborock/api.py

@@ -52,6 +52,7 @@
 from .util import unpack_list
 
 _LOGGER = logging.getLogger(__name__)
+KEEPALIVE = 60
 QUEUE_TIMEOUT = 4
 SPECIAL_COMMANDS = [
     RoborockCommand.GET_MAP_V1,
@@ -89,16 +90,25 @@ def __init__(self, endpoint: str, devices_info: Mapping[str, RoborockDeviceInfo]
         self._endpoint = endpoint
         self._nonce = secrets.token_bytes(16)
         self._waiting_queue: dict[int, RoborockFuture] = {}
-        self._status_listeners: list[Callable[[int, str], None]] = []
+        self._last_device_msg_in = self.time_func()
+        self._last_disconnection = self.time_func()
+        self.keep_alive = KEEPALIVE
 
-    def add_status_listener(self, callback: Callable[[int, str], None]):
-        self._status_listeners.append(callback)
+    @property
+    def time_func(self) -> Callable[[], float]:
+        try:
+            # Use monotonic clock if available
+            time_func = time.monotonic
+        except AttributeError:
+            time_func = time.time
+        return time_func
 
     async def async_disconnect(self) -> Any:
         raise NotImplementedError
 
     def on_message(self, messages: list[RoborockMessage]) -> None:
         try:
+            self._last_device_msg_in = self.time_func()
             for data in messages:
                 protocol = data.protocol
                 if protocol == 102 or protocol == 4:
@@ -142,6 +152,19 @@ def on_message(self, messages: list[RoborockMessage]) -> None:
         except Exception as ex:
             _LOGGER.exception(ex)
 
+    def on_disconnect(self, exc: Optional[Exception]) -> None:
+        self._last_disconnection = self.time_func()
+        _LOGGER.warning("Roborock client disconnected")
+        if exc is not None:
+            _LOGGER.warning(exc)
+
+    def should_keepalive(self) -> bool:
+        now = self.time_func()
+        # noinspection PyUnresolvedReferences
+        if now - self._last_disconnection > self.keep_alive**2 and now - self._last_device_msg_in > self.keep_alive:
+            return False
+        return True
+
     async def _async_response(self, request_id: int, protocol_id: int = 0) -> tuple[Any, VacuumError | None]:
         try:
             queue = RoborockFuture(protocol_id)

roborock/cloud_api.py

@@ -10,15 +10,14 @@
 
 import paho.mqtt.client as mqtt
 
-from .api import SPECIAL_COMMANDS, RoborockClient, md5hex
+from .api import SPECIAL_COMMANDS, RoborockClient, md5hex, KEEPALIVE
 from .containers import RoborockDeviceInfo, UserData
 from .exceptions import CommandVacuumError, RoborockException, VacuumError
 from .roborock_future import RoborockFuture
 from .roborock_message import RoborockMessage, RoborockParser, md5bin
 from .typing import RoborockCommand
 
 _LOGGER = logging.getLogger(__name__)
-MQTT_KEEPALIVE = 60
 CONNECT_REQUEST_ID = 0
 DISCONNECT_REQUEST_ID = 1
 
@@ -49,8 +48,6 @@ def __init__(self, user_data: UserData, devices_info: Mapping[str, RoborockDevic
         self._endpoint = base64.b64encode(md5bin(rriot.k)[8:14]).decode()
         self._waiting_queue: dict[int, RoborockFuture] = {}
         self._mutex = Lock()
-        self._last_device_msg_in = mqtt.time_func()
-        self._last_disconnection = mqtt.time_func()
         self.update_client_id()
 
     def __del__(self) -> None:
@@ -80,19 +77,16 @@ def on_connect(self, *args, **kwargs) -> None:
 
     def on_message(self, *args, **kwargs) -> None:
         _, __, msg = args
-        self._last_device_msg_in = mqtt.time_func()
         device_id = msg.topic.split("/").pop()
         messages, _ = RoborockParser.decode(msg.payload, self.devices_info[device_id].device.local_key)
         super().on_message(messages)
 
     def on_disconnect(self, *args, **kwargs) -> None:
         try:
             _, __, rc, ___ = args
-            self._last_disconnection = mqtt.time_func()
-            message = f"Roborock mqtt client disconnected (rc: {rc})"
+            super().on_disconnect(RoborockException(f"(rc: {rc})"))
             if rc == mqtt.MQTT_ERR_PROTOCOL:
                 self.update_client_id()
-            _LOGGER.warning(message)
             connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
             if connection_queue:
                 connection_queue.resolve((True, None))
@@ -102,18 +96,10 @@ def on_disconnect(self, *args, **kwargs) -> None:
     def update_client_id(self):
         self._client_id = mqtt.base62(uuid.uuid4().int, padding=22)
 
-    def _async_check_keepalive(self) -> None:
-        now = mqtt.time_func()
-        # noinspection PyUnresolvedReferences
-        if (
-            now - self._last_disconnection > self._keepalive**2  # type: ignore[attr-defined]
-            and now - self._last_device_msg_in > self._keepalive  # type: ignore[attr-defined]
-        ):
-            self._ping_t = self._last_device_msg_in
-
     def _check_keepalive(self) -> None:
-        self._async_check_keepalive()
-        # noinspection PyUnresolvedReferences
+        if not self.should_keepalive():
+            self._ping_t = self.time_func() - KEEPALIVE
+            # noinspection PyUnresolvedReferences
         super()._check_keepalive()  # type: ignore[misc]
 
     def sync_stop_loop(self) -> None:
@@ -142,7 +128,7 @@ def sync_connect(self) -> bool:
             if self._mqtt_port is None or self._mqtt_host is None:
                 raise RoborockException("Mqtt information was not entered. Cannot connect.")
             _LOGGER.info("Connecting to mqtt")
-            super().connect_async(host=self._mqtt_host, port=self._mqtt_port, keepalive=MQTT_KEEPALIVE)
+            super().connect_async(host=self._mqtt_host, port=self._mqtt_port, keepalive=KEEPALIVE)
             return True
         return False
 

roborock/exceptions.py

@@ -1,7 +1,7 @@
 """Roborock exceptions."""
 
 
-class RoborockException(BaseException):
+class RoborockException(Exception):
     """Class for Roborock exceptions."""
 
 

roborock/local_api.py

@@ -24,9 +24,7 @@ def __init__(self, devices_info: Mapping[str, RoborockLocalDeviceInfo]):
         self.loop = get_running_loop_or_create_one()
         self.device_listener: dict[str, RoborockSocketListener] = {
             device_id: RoborockSocketListener(
-                device_info.network_info.ip,
-                device_info.device.local_key,
-                self.on_message,
+                device_info.network_info.ip, device_info.device.local_key, self.on_message, self.on_disconnect
             )
             for device_id, device_info in devices_info.items()
         }
@@ -83,6 +81,9 @@ async def send_message(self, device_id: str, roborock_messages: list[RoborockMes
         listener = self.device_listener.get(device_id)
         if listener is None:
             raise RoborockException(f"No device listener for {device_id}")
+        if not self.should_keepalive():
+            listener.disconnect()
+
         _LOGGER.debug(f"Requesting device with {roborock_messages}")
         await listener.send_message(msg)
 
@@ -113,12 +114,14 @@ def __init__(
         ip: str,
         local_key: str,
         on_message: Callable[[list[RoborockMessage]], None],
+        on_disconnect: Callable[[Optional[Exception]], None],
         timeout: float | int = QUEUE_TIMEOUT,
     ):
         self.ip = ip
         self.local_key = local_key
         self.loop = get_running_loop_or_create_one()
         self.on_message = on_message
+        self.on_disconnect = on_disconnect
         self.timeout = timeout
         self.remaining = b""
         self.transport: Transport | None = None
@@ -132,8 +135,8 @@ def data_received(self, message):
         self.remaining = remaining
         self.on_message(parser_msg)
 
-    def connection_lost(self, exc):
-        _LOGGER.debug("The server closed the connection")
+    def connection_lost(self, exc: Optional[Exception]):
+        self.on_disconnect(exc)
 
     def is_connected(self):
         return self.transport and self.transport.is_reading()