Coverage for zenmqtt/mqtt/protocol/outgoing.py: 91%
205 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 21:55 +0000
1import asyncio
2from collections import deque
3from dataclasses import dataclass
4from logging import getLogger
5from typing import Optional, cast
7from zenmqtt.connection import MQTTConnection
8from zenmqtt.exceptions import (
9 ConnectionLostError,
10 NotConnectedError,
11 ServerLimitError,
12 SessionLostError,
13)
14from zenmqtt.mqtt.connect import ConnectionResult
15from zenmqtt.mqtt.limits import DEFAULT_RECEIVE_MAXIMUM
16from zenmqtt.mqtt.protocol.context import ProtocolContext
17from zenmqtt.mqtt.publish import (
18 PubAckResult,
19 PubCompResult,
20 PublishAcknowledgement,
21 PublishProperties,
22 PubRecResult,
23 pack_publish_packet,
24 pack_pubrel_packet,
25)
26from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE, PACKET_IDENTIFIER_NOT_FOUND
27from zenmqtt.session import (
28 MQTTSession,
29 OutgoingMessage,
30 OutgoingMessageState,
31 PacketIdentifier,
32)
34logger = getLogger(__name__)
37class _SendQuota:
38 """
39 Flow control (MQTT 5, section 4.9): limits the number of QoS 1/2
40 publications which are not acknowledged yet by "Receive Maximum".
41 """
43 def __init__(self, limit: int) -> None:
44 self._available = limit
45 self._holders: set[PacketIdentifier] = set()
46 self._waiters: deque[asyncio.Future[None]] = deque()
48 async def acquire(self, packet_identifier: PacketIdentifier) -> bool:
49 """
50 Returns True if the caller had to wait for a free slot.
51 """
52 waited = False
54 if self._available > 0 and not self._waiters:
55 self._available -= 1
56 else:
57 waited = True
58 waiter: asyncio.Future[None] = asyncio.get_running_loop().create_future()
59 self._waiters.append(waiter)
61 try:
62 await waiter
63 except asyncio.CancelledError:
64 # the slot was handed to us right before cancellation
65 if waiter.done() and not waiter.cancelled():
66 self._wake_up_or_return()
67 raise
69 self._holders.add(packet_identifier)
71 return waited
73 def release(self, packet_identifier: PacketIdentifier) -> None:
74 # only messages sent within this connection hold a slot
75 if packet_identifier in self._holders:
76 self._holders.remove(packet_identifier)
77 self._wake_up_or_return()
79 def abort(self, exc: Exception) -> None:
80 while self._waiters: 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true
81 if not (waiter := self._waiters.popleft()).done():
82 waiter.set_exception(exc)
84 def _wake_up_or_return(self) -> None:
85 while self._waiters:
86 if not (waiter := self._waiters.popleft()).done(): 86 ↛ 85line 86 didn't jump to line 85 because the condition on line 86 was always true
87 waiter.set_result(None)
88 return
90 self._available += 1
93@dataclass(slots=True)
94class _PendingPublish:
95 qos: int
96 # there is no future for messages re-sent from the session of a previous run
97 future: Optional[asyncio.Future[PublishAcknowledgement]] = None
98 # loop time when the packet was sent within the current connection
99 sent_at: Optional[float] = None
102def _strip_topic_alias(properties: PublishProperties) -> PublishProperties:
103 # topic aliases don't survive the reconnect
104 return cast(
105 PublishProperties, {k: v for k, v in properties.items() if k != "topic_alias"}
106 )
109class OutgoingFlow:
110 """
111 Outgoing PUBLISH: QoS 1/2 messages are stored in the session until the
112 final acknowledgement and re-sent after reconnect.
113 """
115 def __init__(
116 self,
117 context: ProtocolContext,
118 session: MQTTSession,
119 ) -> None:
120 self._context = context
121 self._session = session
122 # set by close(): publish calls fail instead of waiting for a connect
123 self._closed_with: Optional[Exception] = None
125 self._pending: dict[PacketIdentifier, _PendingPublish] = {}
126 self._send_quota: Optional[_SendQuota] = None
128 # QoS 1/2 publish calls which haven't returned yet
129 self._publishing = 0
130 # set when a message completes or a publish call returns, wakes up
131 # wait_completed()
132 self._progress = asyncio.Event()
134 async def publish(
135 self,
136 topic: str,
137 payload: bytes,
138 qos: int = 0,
139 retain: bool = False,
140 properties: Optional[PublishProperties] = None,
141 ) -> Optional[PublishAcknowledgement]:
142 if qos not in (0, 1, 2):
143 raise ValueError(f"Invalid QoS: {qos}")
145 connection = self._context.ensure_connected()
146 properties = properties or {}
148 # all limits are checked before the packet identifier is acquired
149 self._context.server_limits.check_publish(qos, retain, properties)
151 if qos == 0:
152 packet = pack_publish_packet(
153 0, topic, payload, qos, retain, False, properties
154 )
155 self._context.server_limits.check_packet_size(packet)
157 logger.debug("mqtt_protocol.send_publish_packet")
158 sent_at = self._context.now()
159 await self._context.write(connection, packet)
160 self._context.metrics.on_publish_completed(
161 0, 0, self._context.now() - sent_at
162 )
163 return None
165 self._publishing += 1
167 try:
168 return await self._publish_stored(topic, payload, qos, retain, properties)
169 finally:
170 self._publishing -= 1
171 self._progress.set()
173 async def wait_completed(self) -> None:
174 """
175 Waits until no QoS 1/2 message is in flight: every publish call
176 returned, and every message re-sent from the session is completed or
177 failed.
178 """
179 while self._pending or self._publishing:
180 self._progress.clear()
181 await self._progress.wait()
183 async def _publish_stored(
184 self,
185 topic: str,
186 payload: bytes,
187 qos: int,
188 retain: bool,
189 properties: PublishProperties,
190 ) -> PublishAcknowledgement:
191 packet_identifier = await self._session.acquire_packet_identifier()
193 try:
194 packet = pack_publish_packet(
195 packet_identifier, topic, payload, qos, retain, False, properties
196 )
197 connection = await self._reserve(packet_identifier, packet)
199 await self._session.store_outgoing_message(
200 OutgoingMessage(
201 packet_identifier=packet_identifier,
202 topic=topic,
203 payload=payload,
204 qos=qos,
205 retain=retain,
206 properties=properties,
207 state=OutgoingMessageState.AWAITING_ACK,
208 )
209 )
210 except BaseException:
211 # e.g. the session storage failed or the call was cancelled
212 if self._send_quota: 212 ↛ 215line 212 didn't jump to line 215 because the condition on line 212 was always true
213 self._send_quota.release(packet_identifier)
215 await self._session.release_packet_identifier(packet_identifier)
216 raise
218 # from here the message belongs to the session: if the connection was
219 # lost meanwhile, it'll be re-sent on the next connect
220 return await self._send_stored(packet_identifier, qos, packet, connection)
222 async def on_connected(self, connection_result: ConnectionResult) -> None:
223 """Called after successful CONNACK, before the client is connected."""
224 self._closed_with = None
226 # publish calls wait while the server's "Receive Maximum" is reached
227 self._send_quota = _SendQuota(
228 connection_result.properties.get("receive_maximum", DEFAULT_RECEIVE_MAXIMUM)
229 )
231 await self._resend_pending_messages()
233 def on_session_lost(self) -> None:
234 """The server didn't keep the session, stored messages are discarded."""
235 self._fail_all(SessionLostError())
237 def on_connection_lost(self, exc: ConnectionLostError) -> None:
238 if self._send_quota:
239 # messages waiting for the quota weren't stored nor sent
240 self._send_quota.abort(NotConnectedError())
241 self._send_quota = None
243 # QoS 1/2 messages are re-sent from the session and publish calls get
244 # the acknowledgement of the re-sent message
245 for pending in self._pending.values():
246 pending.sent_at = None
248 def close(self, exc: Exception) -> None:
249 """Nothing will connect anymore, publish calls don't wait."""
250 self._closed_with = exc
251 self._fail_all(exc)
253 async def handle_puback(self, puback_result: PubAckResult) -> None:
254 logger.debug("mqtt_protocol.handle_puback_packet packet:%s", puback_result)
256 await self._complete(puback_result.packet_identifier, puback_result)
258 async def handle_pubrec(self, pubrec_packet: PubRecResult) -> None:
259 connection = self._context.connection
260 assert connection
262 packet_identifier = pubrec_packet.packet_identifier
264 logger.debug("mqtt_protocol.handle_pubrec_packet packet:%s", pubrec_packet)
266 if pubrec_packet.reason_code >= FAILURE_REASON_CODE:
267 # the server rejected the message, the flow ends here
268 await self._complete(packet_identifier, pubrec_packet)
269 return
271 if await self._session.mark_outgoing_message_released(packet_identifier):
272 reason_code = 0x0
273 else:
274 logger.warning(
275 "mqtt_protocol.handle_pubrec_packet.unknown pid:%s", packet_identifier
276 )
277 reason_code = PACKET_IDENTIFIER_NOT_FOUND
279 await self._context.write(
280 connection, pack_pubrel_packet(packet_identifier, reason_code, {})
281 )
283 async def handle_pubcomp(self, pubcomp_packet: PubCompResult) -> None:
284 logger.debug("mqtt_protocol.handle_pubcomp_packet packet:%s", pubcomp_packet)
286 await self._complete(pubcomp_packet.packet_identifier, pubcomp_packet)
288 async def _reserve(
289 self, packet_identifier: PacketIdentifier, packet: bytes
290 ) -> MQTTConnection:
291 """Checks the packet and takes a slot of the send quota."""
292 self._context.server_limits.check_packet_size(packet)
294 self._context.ensure_connected()
296 if self._send_quota: 296 ↛ 300line 296 didn't jump to line 300 because the condition on line 296 was always true
297 await self._acquire_send_quota(self._send_quota, packet_identifier)
299 # the connection could be lost while the call waited for the quota
300 return self._context.ensure_connected()
302 async def _send_stored(
303 self,
304 packet_identifier: PacketIdentifier,
305 qos: int,
306 packet: bytes,
307 connection: MQTTConnection,
308 ) -> PublishAcknowledgement:
309 connected = self._context.connected
311 # e.g. the connection was lost and closed while the message was stored
312 if not connected and self._closed_with:
313 raise self._closed_with
315 future: asyncio.Future[PublishAcknowledgement] = (
316 asyncio.get_running_loop().create_future()
317 )
318 pending = _PendingPublish(qos, future)
319 self._pending[packet_identifier] = pending
321 if connected: 321 ↛ 334line 321 didn't jump to line 334 because the condition on line 321 was always true
322 logger.debug("mqtt_protocol.send_publish_packet pid:%s", packet_identifier)
323 pending.sent_at = self._context.now()
325 try:
326 await self._context.write(connection, packet)
327 except Exception:
328 logger.debug(
329 "mqtt_protocol.send_publish_packet.failed pid:%s, "
330 "it'll be re-sent after reconnect",
331 packet_identifier,
332 )
334 return await future
336 async def _resend_pending_messages(self) -> None:
337 connection = self._context.connection
338 assert connection
340 messages = await self._session.get_pending_outgoing_messages()
342 if messages:
343 self._context.metrics.on_messages_resent(len(messages))
345 for message in messages:
346 packet_identifier = message.packet_identifier
348 if message.state == OutgoingMessageState.AWAITING_ACK:
349 properties = _strip_topic_alias(message.properties)
350 packet = pack_publish_packet(
351 packet_identifier,
352 message.topic,
353 message.payload,
354 message.qos,
355 message.retain,
356 True,
357 properties,
358 )
360 try:
361 self._context.server_limits.check_publish(
362 message.qos, message.retain, properties
363 )
364 self._context.server_limits.check_packet_size(packet)
365 except ServerLimitError as exc:
366 # the server would close the connection on every re-send
367 logger.warning(
368 "mqtt_protocol.resend_pending_message.discard pid:%s reason:%s",
369 packet_identifier,
370 exc,
371 )
372 await self._session.complete_outgoing_message(packet_identifier)
374 if discarded := self._pending.pop(packet_identifier, None): 374 ↛ 375line 374 didn't jump to line 375 because the condition on line 374 was never true
375 self._fail(discarded, exc)
376 self._progress.set()
378 continue
380 if self._send_quota: 380 ↛ 385line 380 didn't jump to line 385 because the condition on line 380 was always true
381 await self._acquire_send_quota(self._send_quota, packet_identifier)
382 else:
383 packet = pack_pubrel_packet(packet_identifier, 0x0, {})
385 logger.debug(
386 "mqtt_protocol.resend_pending_message pid:%s state:%s",
387 packet_identifier,
388 message.state.name,
389 )
391 pending = self._pending.setdefault(
392 packet_identifier, _PendingPublish(message.qos)
393 )
394 pending.sent_at = self._context.now()
396 await self._context.write(connection, packet)
398 async def _complete(
399 self, packet_identifier: PacketIdentifier, result: PublishAcknowledgement
400 ) -> None:
401 await self._session.complete_outgoing_message(packet_identifier)
403 pending = self._pending.pop(packet_identifier, None)
404 self._progress.set()
406 if pending and pending.sent_at is not None: 406 ↛ 411line 406 didn't jump to line 411 because the condition on line 406 was always true
407 self._context.metrics.on_publish_completed(
408 pending.qos, result.reason_code, self._context.now() - pending.sent_at
409 )
411 if self._send_quota: 411 ↛ 414line 411 didn't jump to line 414 because the condition on line 411 was always true
412 self._send_quota.release(packet_identifier)
414 if pending and pending.future and not pending.future.done():
415 pending.future.set_result(result)
417 async def _acquire_send_quota(
418 self, send_quota: _SendQuota, packet_identifier: PacketIdentifier
419 ) -> None:
420 started_at = self._context.now()
422 if await send_quota.acquire(packet_identifier):
423 self._context.metrics.on_send_quota_wait(self._context.now() - started_at)
425 def _fail_all(self, exc: Exception) -> None:
426 pending, self._pending = self._pending, {}
427 self._progress.set()
429 for publish in pending.values():
430 self._fail(publish, exc)
432 @staticmethod
433 def _fail(pending: _PendingPublish, exc: Exception) -> None:
434 if pending.future and not pending.future.done(): 434 ↛ exitline 434 didn't return from function '_fail' because the condition on line 434 was always true
435 pending.future.set_exception(exc)