Coverage for zenmqtt/client.py: 95%
241 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 21:55 +0000
1import asyncio
2import inspect
3from contextlib import AbstractContextManager, nullcontext
4from dataclasses import dataclass
5from logging import getLogger
6from ssl import SSLContext
7from types import TracebackType
8from typing import Any, Awaitable, Callable, Final, Optional, Sequence, TypeVar
10from zenmqtt.connection import create_connection
11from zenmqtt.exceptions import ConnectionLostError, NotConnectedError
12from zenmqtt.instrumentation.metrics import MetricsCollector, guard_metrics
13from zenmqtt.instrumentation.tracing import Tracing, guard_tracing
14from zenmqtt.mqtt.connect import (
15 SESSION_PRESENT_FLAG,
16 ConnectionResult,
17 ConnectProperties,
18 DisconnectProperties,
19 ServerDisconnect,
20 WillMessage,
21)
22from zenmqtt.mqtt.protocol import MQTTProtocol
23from zenmqtt.mqtt.publish import (
24 MQTTMessage,
25 PublishAcknowledgement,
26 PublishProperties,
27)
28from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE
29from zenmqtt.mqtt.subscribe import (
30 SubscribeResult,
31 SubscriptionProperties,
32 SubscriptionRequest,
33 UnsubscribeProperties,
34 UnsubscribeResult,
35)
36from zenmqtt.reconnect import DEFAULT_RECONNECT_POLICY, Reconnector, ReconnectPolicy
37from zenmqtt.session import MQTTSession, build_default_session
38from zenmqtt.subscriptions import SubscriptionRegistry
40logger = getLogger(__name__)
42DEFAULT_KEEP_ALIVE: Final[int] = 60
43# seconds which `async with MQTTClient(...)` waits on exit for outgoing QoS 1/2
44# messages in flight, see MQTTClient.__aexit__
45DEFAULT_DISCONNECT_TIMEOUT: Final[float] = 10.0
48T = TypeVar("T")
50OnConnect = Callable[[ConnectionResult], Any]
51OnDisconnect = Callable[[ConnectionLostError], Any]
54class AsyncMessageIterator:
55 def __init__(
56 self, next_message: Callable[[], Awaitable[Optional[MQTTMessage]]]
57 ) -> None:
58 self._next_message = next_message
60 def __aiter__(self) -> "AsyncMessageIterator":
61 return self
63 async def __anext__(self) -> MQTTMessage:
64 value = await self._next_message()
66 if value is None:
67 raise StopAsyncIteration
69 return value
72@dataclass(frozen=True, slots=True)
73class _ConnectOptions:
74 url: str
75 username: Optional[str]
76 password: Optional[str]
77 keepalive: int
78 properties: Optional[ConnectProperties]
79 will: Optional[WillMessage]
80 ssl: Optional[SSLContext]
83class MQTTClient:
84 def __init__(
85 self,
86 client_id: str,
87 session: Optional[MQTTSession] = None,
88 metrics: Optional[MetricsCollector] = None,
89 reconnect: Optional[ReconnectPolicy] = DEFAULT_RECONNECT_POLICY,
90 tracing: Optional[Tracing] = None,
91 ) -> None:
92 """
93 :param session: storage of in-flight QoS 1/2 messages, in-memory by default
94 :param metrics: receives events of the client to build metrics,
95 e.g. zenmqtt.instrumentation.opentelemetry.OpenTelemetryMetrics
96 :param reconnect: automatic reconnect after the connection is lost,
97 None disables it; see connect() for details
98 :param tracing: spans of publish calls and of the processing of
99 messages, e.g. zenmqtt.instrumentation.opentelemetry.OpenTelemetryTracing;
100 the trace context is sent in user properties
101 """
102 # hooks are guarded once here, the protocol keeps the guarded collector;
103 # instrumentation never changes the behaviour of the client
104 self._metrics = guard_metrics(metrics)
105 self._tracing = guard_tracing(tracing)
106 self._reconnect_policy = reconnect
108 self._session = session or build_default_session()
109 self._protocol = MQTTProtocol(self._session, self._metrics)
111 self._client_id = client_id
113 self._connect_options: Optional[_ConnectOptions] = None
114 # a task waits for the next message, see messages
115 self._reading = False
117 self._subscriptions = SubscriptionRegistry()
118 self._reconnector: Optional[Reconnector] = None
119 # without reconnect: closes the protocol when the connection is lost
120 self._watcher: Optional[asyncio.Task[None]] = None
122 self._connected = False
123 # True until connect and after the client stops for good
124 self._stopped = True
125 # disconnect(timeout=...) waits for messages in flight, new operations
126 # are rejected
127 self._draining = False
128 self._state_changed = asyncio.Event()
130 # called on every successful connect and on every lost connection,
131 # may be sync or async functions; one callback each: assigning a new
132 # one replaces the previous one
133 self.on_connect: Optional[OnConnect] = None
134 self.on_disconnect: Optional[OnDisconnect] = None
136 async def __aenter__(self) -> "MQTTClient":
137 """
138 `async with MQTTClient(...) as client:` makes sure the client is
139 disconnected when the block ends; connect() is called inside the block.
140 """
141 return self
143 async def __aexit__(
144 self,
145 exc_type: Optional[type[BaseException]],
146 exc: Optional[BaseException],
147 traceback: Optional[TracebackType],
148 ) -> None:
149 """
150 Disconnects the client if it was connected (also while it reconnects).
151 On a normal exit the shutdown is graceful: the client waits up to
152 DEFAULT_DISCONNECT_TIMEOUT seconds for outgoing QoS 1/2 messages in
153 flight. If the block raised (or was cancelled), the client disconnects
154 at once and doesn't hide the exception; messages in flight stay in the
155 session. For another timeout call disconnect(timeout=...) in the block.
156 """
157 if self._stopped:
158 return
160 if exc_type is None:
161 await self.disconnect(timeout=DEFAULT_DISCONNECT_TIMEOUT)
162 return
164 try:
165 await self.disconnect()
166 except Exception as error:
167 logger.error("mqtt_client.disconnect.error", exc_info=error)
169 @property
170 def client_id(self) -> str:
171 """
172 The client identifier; if the server assigned one (the client id was
173 empty), the assigned identifier, which is used for the next connects.
174 """
175 return self._client_id
177 async def connect(
178 self,
179 url: str,
180 *,
181 username: Optional[str] = None,
182 password: Optional[str] = None,
183 clean_start: bool = False,
184 keepalive: int = DEFAULT_KEEP_ALIVE,
185 properties: Optional[ConnectProperties] = None,
186 will: Optional[WillMessage] = None,
187 ssl: Optional[SSLContext] = None,
188 ) -> ConnectionResult:
189 """
190 If the connection can't be established, the error is raised: automatic
191 reconnect starts only after a successful connect. Then the client
192 reconnects with the same options, but without Clean Start, so the
193 session of the server continues; QoS 1/2 messages are re-sent and lost
194 subscriptions are restored.
196 :param url: tcp://host[:port] or mqtts://host[:port] for TLS, the default
197 port is 1883 for tcp:// and 8883 for mqtts://
198 :param username: MQTT 5 allows a password without a username, both may
199 be empty strings; reconnects use the same credentials
200 :param clean_start: discard the session on the server and the client;
201 to re-send QoS 1/2 messages after a reconnect use False together
202 with properties={"session_expiry_interval": <seconds>}, otherwise
203 the server drops the session when the connection is closed
204 :param keepalive: seconds between control packets sent by the client,
205 the client sends PINGREQ when idle; 0 disables keep alive
206 :param will: the message the server publishes if the client goes away
207 without DISCONNECT
208 :param ssl: context for mqtts://, e.g. with a custom CA or a client
209 certificate; by default the server certificate is verified with
210 the system CA certificates
211 """
212 if self._reconnector and self._reconnector.running: 212 ↛ 213line 212 didn't jump to line 213 because the condition on line 212 was never true
213 raise RuntimeError("The client is already connected")
215 self._connect_options = _ConnectOptions(
216 url, username, password, keepalive, properties, will, ssl
217 )
219 if self._reconnect_policy and not (properties or {}).get(
220 "session_expiry_interval"
221 ):
222 logger.warning(
223 "mqtt_client.connect: reconnect is enabled, but the server drops "
224 "the session when the connection is lost; set "
225 "session_expiry_interval > 0 to keep QoS 1/2 messages and "
226 "subscriptions"
227 )
229 connack = await self._connect(clean_start)
231 if connack.reason_code >= FAILURE_REASON_CODE: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 return connack
234 self._stopped = False
235 self._set_connected(True)
236 await self._notify(self.on_connect, connack)
238 if self._reconnect_policy:
239 self._reconnector = Reconnector(
240 self._reconnect_policy,
241 self._protocol,
242 self._metrics,
243 connect=lambda: self._connect(clean_start=False),
244 connection_lost=self._on_connection_lost,
245 reconnected=self._on_reconnected,
246 gave_up=self._on_gave_up,
247 )
248 self._reconnector.start()
249 else:
250 self._watcher = asyncio.create_task(self._watch_connection())
252 return connack
254 async def disconnect(
255 self,
256 reason_code: int = 0,
257 properties: Optional[DisconnectProperties] = None,
258 *,
259 timeout: Optional[float] = None,
260 ) -> int:
261 """
262 Stops automatic reconnect; pending publish calls fail with
263 ConnectionLostError, the messages iterator ends.
265 With a timeout the shutdown is graceful: new publish, subscribe and
266 unsubscribe calls fail with NotConnectedError, and the client waits up
267 to `timeout` seconds until the server acknowledged every outgoing QoS
268 1/2 message in flight (PUBACK, PUBCOMP); then it sends DISCONNECT. If
269 the connection is lost meanwhile, the client keeps reconnecting (with
270 automatic reconnect) and the re-sent messages are waited for too. QoS 0
271 messages are sent at once, so they aren't waited for.
273 :param reason_code: reason code of DISCONNECT, e.g. 0x04 "Disconnect with
274 Will Message" asks the server to publish the will message
275 :param timeout: seconds to wait for messages in flight; None
276 disconnects at once
277 :returns: the number of outgoing QoS 1/2 messages which the server
278 didn't acknowledge; they stay in the session and are re-sent on the
279 next connect with this session
280 """
281 was_stopped = self._stopped
283 if timeout is not None and not was_stopped:
284 await self._drain(timeout)
286 self._stopped = True
288 if reconnector := self._reconnector:
289 self._reconnector = None
290 await reconnector.stop()
292 await self._protocol.disconnect(reason_code=reason_code, properties=properties)
294 self._set_connected(False)
296 if not was_stopped:
297 await self._protocol.close(ConnectionLostError())
299 self._draining = False
300 left = len(await self._session.get_pending_outgoing_messages())
302 if left and timeout is not None:
303 logger.warning(
304 "mqtt_client.disconnect: %s outgoing QoS 1/2 messages aren't "
305 "acknowledged by the server within %s s, they stay in the session",
306 left,
307 timeout,
308 )
310 return left
312 async def _drain(self, timeout: float) -> None:
313 self._draining = True
315 try:
316 await asyncio.wait_for(self._protocol.wait_outgoing_completed(), timeout)
317 except asyncio.TimeoutError:
318 pass
320 @property
321 def is_connected(self) -> bool:
322 # the protocol notices a lost connection before the reconnector
323 return self._connected and self._protocol.is_connected
325 async def wait_connected(self) -> None:
326 """
327 Waits while the client reconnects.
329 :raises NotConnectedError: the client isn't connected and doesn't
330 reconnect (not connected yet, disconnected or gave up)
331 """
332 while not self.is_connected:
333 if self._stopped:
334 raise NotConnectedError()
336 await self._state_changed.wait()
338 @property
339 def server_disconnect(self) -> Optional[ServerDisconnect]:
340 """
341 DISCONNECT packet (reason code and properties) if the server closed
342 the last connection with it, otherwise None; reset on connect.
343 """
344 return self._protocol.server_disconnect
346 async def publish(
347 self,
348 topic: str,
349 message: bytes,
350 qos: int = 0,
351 retain: bool = False,
352 properties: Optional[PublishProperties] = None,
353 ) -> Optional[PublishAcknowledgement]:
354 """
355 Returns None for QoS 0, otherwise the final acknowledgement packet.
357 With automatic reconnect, the call waits while the client reconnects,
358 and QoS 1/2 calls wait for the acknowledgement of the re-sent message;
359 use asyncio.wait_for to limit the time. A cancelled QoS 1/2 call
360 doesn't cancel the delivery: the message stays in the session.
362 :raises ValueError: QoS isn't 0, 1 or 2
363 :raises ServerLimitError: the message breaks a limit of the server from
364 CONNACK: QoSNotSupportedError, OutgoingPacketTooLargeError or
365 FeatureNotSupportedError (retain, topic alias); nothing is sent
366 :raises NotConnectedError: the client isn't connected and doesn't
367 reconnect
368 :raises SessionLostError: the server didn't keep the session after
369 reconnect, the message was discarded
370 :raises ConnectionLostError: without automatic reconnect: the connection
371 was lost before the acknowledgement, the message will be re-sent on
372 the next connect only if the server keeps the session, see
373 connect(clean_start); with automatic reconnect: the client
374 stopped before the acknowledgement
375 """
376 if self._tracing is None:
377 return await self._when_connected(
378 lambda: self._protocol.publish(
379 topic, message, qos=qos, retain=retain, properties=properties
380 )
381 )
383 trace = self._tracing.start_publish(
384 self.client_id, topic, message, qos, retain, properties or {}
385 )
387 try:
388 result = await self._when_connected(
389 lambda: self._protocol.publish(
390 topic, message, qos=qos, retain=retain, properties=trace.properties
391 )
392 )
393 except BaseException as exc:
394 trace.end(None, exc)
395 raise
397 trace.end(result, None)
398 return result
400 async def ping(self) -> None:
401 """
402 Sends PINGREQ and waits for PINGRESP.
404 :raises NotConnectedError: the client isn't connected
405 :raises ConnectionLostError: the connection was lost before PINGRESP
406 """
407 await self._protocol.ping()
409 async def subscribe(
410 self,
411 topics: Sequence[SubscriptionRequest],
412 properties: Optional[SubscriptionProperties] = None,
413 ) -> SubscribeResult:
414 """
415 With automatic reconnect, the call waits while the client reconnects
416 and is repeated if the connection is lost before SUBACK.
418 :param topics: (topic, qos) pairs, or zenmqtt.mqtt.subscribe.Subscription to set
419 MQTT 5 subscription options (no_local, retain_as_published,
420 retain_handling)
421 :raises ValueError: invalid QoS or subscription options
422 """
423 subscription_properties = properties or {}
425 result = await self._when_connected(
426 lambda: self._protocol.subscribe(topics, subscription_properties),
427 repeat=True,
428 )
430 self._subscriptions.add(topics, subscription_properties, result.reason_codes)
432 return result
434 async def unsubscribe(
435 self,
436 topics: Sequence[str],
437 properties: Optional[UnsubscribeProperties] = None,
438 ) -> UnsubscribeResult:
439 """
440 With automatic reconnect, the call waits while the client reconnects
441 and is repeated if the connection is lost before UNSUBACK.
442 """
443 unsubscribe_properties = properties or {}
445 result = await self._when_connected(
446 lambda: self._protocol.unsubscribe(topics, unsubscribe_properties),
447 repeat=True,
448 )
450 self._subscriptions.remove(topics, result.reason_codes)
452 return result
454 async def ack(self, message: MQTTMessage, reason_code: int = 0) -> None:
455 """
456 Acknowledges a processed message: sends PUBACK for QoS 1 and PUBREC for
457 QoS 2, does nothing for QoS 0. Every QoS 1/2 message from
458 client.messages must be acknowledged.
460 Acks are sent in the order the messages were received (MQTT 5, 4.6),
461 so an ack waits for acks of earlier messages: one message which isn't
462 acknowledged holds back the acks of all later ones. Not acknowledged
463 messages count to "Receive Maximum": when it's reached, the server
464 stops sending QoS 1/2 messages. A warning is logged when a message
465 holds back later acks for 30 seconds. If the connection is lost before
466 the ack, the ack is ignored and the server sends the message again
467 after reconnect.
469 :param reason_code: 0, or >= 0x80 to reject the message, e.g. 0x80
470 (Unspecified error) or 0x99 (Payload format invalid); the server
471 doesn't send a rejected message again
472 """
473 await self._protocol.ack(message, reason_code)
475 def trace(self, message: MQTTMessage) -> AbstractContextManager[None]:
476 """
477 Wraps the processing of a received message for tracing: with tracing,
478 the block runs in a span which continues the trace of the publisher;
479 without tracing it does nothing.
481 async for message in client.messages:
482 with client.trace(message):
483 await handle(message)
484 await client.ack(message)
485 """
486 if self._tracing is None:
487 return nullcontext()
489 return self._tracing.trace(self.client_id, message)
491 @property
492 def messages(self) -> AsyncMessageIterator:
493 """
494 Incoming messages; with automatic reconnect the iteration continues
495 across reconnects and ends when the client stops. QoS 1/2 messages must
496 be acknowledged by ack() after they are processed.
498 Messages wait in a buffer while the application reads previous ones,
499 so control packets (PINGRESP, PUBACK, ...) are handled anyway. The
500 number of QoS 1/2 messages which aren't acknowledged is limited by
501 "receive_maximum" of connect properties; QoS 0 messages aren't
502 limited. While 1000 or more messages wait, a warning is logged and
503 metrics get on_messages_buffered events.
505 Messages are read by one task: if another task waits for a message at
506 the same time, it gets RuntimeError. To process messages concurrently,
507 read them in one task and hand them to workers.
508 """
509 return AsyncMessageIterator(self._next_message)
511 async def _next_message(self) -> Optional[MQTTMessage]:
512 if self._reading:
513 raise RuntimeError(
514 "client.messages is read by another task; read messages in one "
515 "task and hand them to workers to process them concurrently"
516 )
518 self._reading = True
520 try:
521 return await self._protocol.next_message()
522 finally:
523 self._reading = False
525 async def _connect(self, clean_start: bool) -> ConnectionResult:
526 options = self._connect_options
527 assert options
529 connection = await create_connection(options.url, ssl=options.ssl)
530 self._protocol.set_connection(connection)
532 try:
533 connack = await self._protocol.connect(
534 self._client_id,
535 options.username,
536 options.password,
537 clean_start,
538 options.keepalive,
539 options.properties,
540 options.will,
541 )
542 except BaseException:
543 # e.g. timeout or cancellation while waiting for CONNACK
544 if not connection.is_closing(): 544 ↛ 546line 544 didn't jump to line 546 because the condition on line 544 was always true
545 await connection.disconnect()
546 raise
548 if connack.reason_code >= FAILURE_REASON_CODE:
549 # the server closes the connection anyway
550 if not connection.is_closing(): 550 ↛ 553line 550 didn't jump to line 553 because the condition on line 550 was always true
551 await connection.disconnect()
553 return connack
555 # the server assigns an identifier if the client id is empty, it must be
556 # used for the following connects to resume the session
557 if assigned := connack.properties.get("assigned_client_identifier"):
558 self._client_id = assigned
560 return connack
562 async def _watch_connection(self) -> None:
563 await self._protocol.wait_closed()
565 # disconnect() closes the protocol itself
566 if self._stopped:
567 return
569 lost = ConnectionLostError(self._protocol.server_disconnect)
570 logger.warning("mqtt_client.connection_lost: %s", lost)
572 self._stopped = True
573 await self._protocol.close(lost)
574 await self._on_connection_lost(lost)
576 async def _on_connection_lost(self, lost: ConnectionLostError) -> None:
577 self._set_connected(False)
578 await self._notify(self.on_disconnect, lost)
580 async def _on_reconnected(self, connack: ConnectionResult) -> None:
581 if not connack.flags & SESSION_PRESENT_FLAG:
582 await self._subscriptions.restore(self._protocol.subscribe)
584 self._set_connected(True)
585 await self._notify(self.on_connect, connack)
587 async def _on_gave_up(self, exc: Exception) -> None:
588 self._stopped = True
589 self._reconnector = None
590 self._set_connected(False)
592 await self._protocol.close(exc)
594 async def _when_connected(
595 self, operation: Callable[[], Awaitable[T]], repeat: bool = False
596 ) -> T:
597 """
598 Without automatic reconnect it just runs the operation. Otherwise waits
599 while the client reconnects; the operation is repeated if the client
600 lost the connection before sending it, or with repeat=True (for
601 idempotent operations) before the acknowledgement.
602 """
603 if self._draining:
604 raise NotConnectedError("The client is disconnecting")
606 if not self._reconnect_policy:
607 return await operation()
609 while True:
610 await self.wait_connected()
612 # the client started to disconnect while this call waited
613 if self._draining: 613 ↛ 614line 613 didn't jump to line 614 because the condition on line 613 was never true
614 raise NotConnectedError("The client is disconnecting")
616 try:
617 return await operation()
618 except NotConnectedError:
619 pass
620 except ConnectionLostError:
621 if not repeat:
622 raise
624 if self._stopped: 624 ↛ 625line 624 didn't jump to line 625 because the condition on line 624 was never true
625 raise NotConnectedError()
627 def _set_connected(self, connected: bool) -> None:
628 self._connected = connected
630 # wake up everybody who waits for a change
631 state_changed, self._state_changed = self._state_changed, asyncio.Event()
632 state_changed.set()
634 @staticmethod
635 async def _notify(callback: Optional[Callable[[Any], Any]], argument: Any) -> None:
636 if callback is None:
637 return
639 try:
640 result = callback(argument)
642 if inspect.isawaitable(result):
643 await result
644 except Exception as exc:
645 logger.error("mqtt_client.callback.error", exc_info=exc)