1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
|
from __future__ import annotations
import asyncio
import logging
import threading
from abc import ABC
from asyncio import Lock
from typing import Any
import paho.mqtt.client as mqtt
from paho.mqtt.enums import MQTTErrorCode
# Mypy is not seeing this for some reason. It wants me to use the depreciated ReasonCodes
from paho.mqtt.reasoncodes import ReasonCode # type: ignore
from .api import KEEPALIVE, RoborockClient
from .containers import DeviceData, UserData
from .exceptions import RoborockException, VacuumError
from .protocol import (
Decoder,
Encoder,
create_mqtt_decoder,
create_mqtt_encoder,
create_mqtt_params,
)
from .roborock_future import RoborockFuture
_LOGGER = logging.getLogger(__name__)
CONNECT_REQUEST_ID = 0
DISCONNECT_REQUEST_ID = 1
class _Mqtt(mqtt.Client):
"""Internal MQTT client.
This is a subclass of the Paho MQTT client that adds some additional functionality
for error cases where things get stuck.
"""
_thread: threading.Thread
def __init__(self) -> None:
"""Initialize the MQTT client."""
super().__init__(protocol=mqtt.MQTTv5)
def maybe_restart_loop(self) -> None:
"""Ensure that the MQTT loop is running in case it previously exited."""
if not self._thread or not self._thread.is_alive():
if self._thread:
_LOGGER.info("Stopping mqtt loop")
super().loop_stop()
_LOGGER.info("Starting mqtt loop")
super().loop_start()
class RoborockMqttClient(RoborockClient, ABC):
"""Roborock MQTT client base class."""
def __init__(self, user_data: UserData, device_info: DeviceData) -> None:
"""Initialize the Roborock MQTT client."""
rriot = user_data.rriot
if rriot is None:
raise RoborockException("Got no rriot data from user_data")
RoborockClient.__init__(self, device_info)
mqtt_params = create_mqtt_params(rriot)
self._mqtt_user = rriot.u
self._hashed_user = mqtt_params.username
self._mqtt_host = mqtt_params.host
self._mqtt_port = mqtt_params.port
self._mqtt_client = _Mqtt()
self._mqtt_client.on_connect = self._mqtt_on_connect
self._mqtt_client.on_message = self._mqtt_on_message
# Due to the incorrect ReasonCode, it is confused by typing
self._mqtt_client.on_disconnect = self._mqtt_on_disconnect # type: ignore
if mqtt_params.tls:
self._mqtt_client.tls_set()
self._mqtt_client.username_pw_set(mqtt_params.username, mqtt_params.password)
self._waiting_queue: dict[int, RoborockFuture] = {}
self._mutex = Lock()
self._decoder: Decoder = create_mqtt_decoder(device_info.device.local_key)
self._encoder: Encoder = create_mqtt_encoder(device_info.device.local_key)
self.received_message_since_last_disconnect = False
self._topic = f"rr/m/o/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}"
def _mqtt_on_connect(
self,
client: mqtt.Client,
userdata: object,
flags: dict[str, int],
rc: ReasonCode,
properties: mqtt.Properties | None = None,
):
connection_queue = self._waiting_queue.get(CONNECT_REQUEST_ID)
if rc.is_failure:
message = f"Failed to connect ({rc})"
self._logger.error(message)
if connection_queue:
connection_queue.set_exception(VacuumError(message))
else:
self._logger.debug("Failed to notify connect future, not in queue")
return
self._logger.info(f"Connected to mqtt {self._mqtt_host}:{self._mqtt_port}")
(result, mid) = self._mqtt_client.subscribe(self._topic)
if result != 0:
message = f"Failed to subscribe ({str(rc)})"
self._logger.error(message)
if connection_queue:
connection_queue.set_exception(VacuumError(message))
return
self._logger.info(f"Subscribed to topic {self._topic}")
if connection_queue:
connection_queue.set_result(True)
def _mqtt_on_message(self, *args, **kwargs):
self.received_message_since_last_disconnect = True
client, __, msg = args
try:
messages = self._decoder(msg.payload)
super().on_message_received(messages)
except Exception as ex:
self._logger.exception(ex)
def _mqtt_on_disconnect(
self,
client: mqtt.Client,
data: object,
flags: dict[str, int],
rc: ReasonCode | None,
properties: mqtt.Properties | None = None,
):
try:
exc = RoborockException(str(rc)) if rc is not None and rc.is_failure else None
super().on_connection_lost(exc)
connection_queue = self._waiting_queue.get(DISCONNECT_REQUEST_ID)
if connection_queue:
connection_queue.set_result(True)
except Exception as ex:
self._logger.exception(ex)
def is_connected(self) -> bool:
"""Check if the mqtt client is connected."""
return self._mqtt_client.is_connected()
def _sync_disconnect(self) -> Any:
if not self.is_connected():
return None
self._logger.info("Disconnecting from mqtt")
disconnected_future = self._async_response(DISCONNECT_REQUEST_ID)
rc = self._mqtt_client.disconnect()
if rc == mqtt.MQTT_ERR_NO_CONN:
disconnected_future.cancel()
return None
if rc != mqtt.MQTT_ERR_SUCCESS:
disconnected_future.cancel()
raise RoborockException(f"Failed to disconnect ({str(rc)})")
return disconnected_future
def _sync_connect(self) -> Any:
if self.is_connected():
self._mqtt_client.maybe_restart_loop()
return None
if self._mqtt_port is None or self._mqtt_host is None:
raise RoborockException("Mqtt information was not entered. Cannot connect.")
self._logger.debug("Connecting to mqtt")
connected_future = self._async_response(CONNECT_REQUEST_ID)
self._mqtt_client.connect(host=self._mqtt_host, port=self._mqtt_port, keepalive=KEEPALIVE)
self._mqtt_client.maybe_restart_loop()
return connected_future
async def async_disconnect(self) -> None:
async with self._mutex:
if disconnected_future := self._sync_disconnect():
# There are no errors set on this future
await disconnected_future
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._mqtt_client.loop_stop)
async def async_connect(self) -> None:
async with self._mutex:
if connected_future := self._sync_connect():
try:
await connected_future
except VacuumError as err:
raise RoborockException(err) from err
def _send_msg_raw(self, msg: bytes) -> None:
info = self._mqtt_client.publish(
f"rr/m/i/{self._mqtt_user}/{self._hashed_user}/{self.device_info.device.duid}", msg
)
if info.rc != mqtt.MQTT_ERR_SUCCESS:
raise RoborockException(f"Failed to publish ({mqtt.error_string(info.rc)})")
async def _unsubscribe(self) -> MQTTErrorCode:
"""Unsubscribe from the topic."""
loop = asyncio.get_running_loop()
(result, mid) = await loop.run_in_executor(None, self._mqtt_client.unsubscribe, self._topic)
if result != 0:
message = f"Failed to unsubscribe ({mqtt.error_string(result)})"
self._logger.error(message)
else:
self._logger.info(f"Unsubscribed from topic {self._topic}")
return result
async def _subscribe(self) -> MQTTErrorCode:
"""Subscribe to the topic."""
loop = asyncio.get_running_loop()
(result, mid) = await loop.run_in_executor(None, self._mqtt_client.subscribe, self._topic)
if result != 0:
message = f"Failed to subscribe ({mqtt.error_string(result)})"
self._logger.error(message)
else:
self._logger.info(f"Subscribed to topic {self._topic}")
return result
async def _reconnect(self) -> None:
"""Reconnect to the MQTT broker."""
await self.async_disconnect()
await self.async_connect()
async def _validate_connection(self) -> None:
"""Override the default validate connection to try to re-subscribe rather than disconnect.
When something seems to be wrong with our connection, we should follow the following steps:
1. Try to unsubscribe and resubscribe from the topic.
2. If we don't end up getting a message, we should completely disconnect and reconnect to the MQTT broker.
3. We will continue to try to disconnect and reconnect until we get a message.
4. If we get a message, the next time connection is lost, We will go back to step 1.
"""
# If we should no longer keep the current connection alive...
if not self.should_keepalive():
self._logger.info("Resetting Roborock connection due to keepalive timeout")
if not self.received_message_since_last_disconnect:
# If we have already tried to unsub and resub, and we are still in this state,
# we should try to reconnect.
return await self._reconnect()
try:
# Mark that we have tried to unsubscribe and resubscribe
self.received_message_since_last_disconnect = False
if await self._unsubscribe() != 0:
# If we fail to unsubscribe, reconnect to the broker
return await self._reconnect()
if await self._subscribe() != 0:
# If we fail to subscribe, reconnected to the broker.
return await self._reconnect()
except Exception: # noqa
# If we get any errors at all, we should just reconnect.
return await self._reconnect()
# Call connect to make sure everything is still in a good state.
await self.async_connect()
|
