Coverage for zenmqtt/mqtt/protocol/incoming.py: 99%

1import asyncio 

2import dataclasses 

3from collections import OrderedDict, deque 

4from dataclasses import dataclass, field 

5from functools import partial 

6from logging import getLogger 

7from typing import Optional, cast 

8 

9from zenmqtt.exceptions import ( 

10 ProtocolError, 

11 ReceiveMaximumExceededError, 

12 TopicAliasInvalidError, 

13) 

14from zenmqtt.instrumentation.monitors import AckBlockedMonitor, BufferMonitor 

15from zenmqtt.mqtt.connect import ConnectProperties 

16from zenmqtt.mqtt.limits import DEFAULT_RECEIVE_MAXIMUM 

17from zenmqtt.mqtt.protocol.context import ProtocolContext 

18from zenmqtt.mqtt.publish import ( 

19 MQTTMessage, 

20 PubRelResult, 

21 pack_puback_packet, 

22 pack_pubcomp_packet, 

23 pack_pubrec_packet, 

24) 

25from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE, PACKET_IDENTIFIER_NOT_FOUND 

26from zenmqtt.session import MQTTSession, PacketIdentifier 

27 

28logger = getLogger(__name__) 

29 

30 

31@dataclass(slots=True) 

32class _PendingAck: 

33 message: MQTTMessage 

34 # loop time 

35 received_at: float 

36 acked: bool = False 

37 reason_code: int = 0 

38 

39 

40@dataclass(slots=True) 

41class _ConnectionState: 

42 """State of incoming messages which doesn't survive the reconnect.""" 

43 

44 topic_aliases: dict[int, str] = field(default_factory=dict) 

45 # QoS 1/2 messages which aren't acknowledged yet; QoS 2 stays here until 

46 # PUBREL, the number is limited by "Receive Maximum" of the client 

47 inflight: set[PacketIdentifier] = field(default_factory=set) 

48 # QoS 1/2 messages in the order of receiving: PUBACK and PUBREC must be 

49 # sent in this order (MQTT 5, 4.6). OrderedDict: getting the first item of 

50 # a dict after deletions from its beginning gets slower with every deleted 

51 # item, of an OrderedDict it takes constant time 

52 pending_acks: OrderedDict[PacketIdentifier, _PendingAck] = field( 

53 default_factory=OrderedDict 

54 ) 

55 

56 

57class IncomingFlow: 

58 """ 

59 Incoming PUBLISH. The read loop never waits for the application: messages 

60 wait in a buffer which the application reads by next_message(), so control 

61 packets (PINGRESP, PUBACK, ...) are handled for a slow application. 

62 

63 QoS 1/2 messages are acknowledged by ack(), so their number in the buffer 

64 is limited by "Receive Maximum"; QoS 0 messages aren't limited. A big 

65 buffer is reported by BufferMonitor. 

66 """ 

67 

68 def __init__(self, context: ProtocolContext, session: MQTTSession) -> None: 

69 self._context = context 

70 self._session = session 

71 

72 # limits sent by the client in CONNECT 

73 self._receive_maximum = DEFAULT_RECEIVE_MAXIMUM 

74 self._topic_alias_maximum = 0 

75 

76 self._state = _ConnectionState() 

77 self._ack_lock = asyncio.Lock() 

78 self._ack_monitor = AckBlockedMonitor() 

79 

80 # messages which wait for the application; None marks the end of 

81 # messages 

82 self._buffer: deque[Optional[MQTTMessage]] = deque() 

83 self._buffer_ready = asyncio.Event() 

84 self.buffer_monitor = BufferMonitor(context.metrics) 

85 

86 @property 

87 def inflight(self) -> set[PacketIdentifier]: 

88 return self._state.inflight 

89 

90 def on_new_connection(self, properties: ConnectProperties) -> None: 

91 """ 

92 Topic aliases and the receive quota don't survive the reconnect; acks 

93 of messages from the previous connection are ignored, the server 

94 re-sends them. 

95 """ 

96 self._receive_maximum = properties.get( 

97 "receive_maximum", DEFAULT_RECEIVE_MAXIMUM 

98 ) 

99 self._topic_alias_maximum = properties.get("topic_alias_maximum", 0) 

100 

101 self._state = _ConnectionState() 

102 

103 def on_connection_lost(self) -> None: 

104 # not acknowledged messages are re-sent by the server, QoS 0 messages 

105 # are delivered 

106 self._buffer = deque( 

107 message for message in self._buffer if message is None or not message.qos 

108 ) 

109 

110 def close(self) -> None: 

111 """The client stops for good: the end of messages.""" 

112 self._enqueue(None) 

113 

114 async def next_message(self) -> Optional[MQTTMessage]: 

115 """ 

116 Waits for the next message, None is the end of messages. Meant for one 

117 reader: the end of messages goes to one of the waiting readers only. 

118 """ 

119 while not self._buffer: 

120 self._buffer_ready.clear() 

121 await self._buffer_ready.wait() 

122 

123 message = self._buffer.popleft() 

124 self.buffer_monitor.update(len(self._buffer)) 

125 

126 return message 

127 

128 async def handle_publish(self, packet: MQTTMessage) -> None: 

129 message = self._resolve_topic_alias(packet) 

130 

131 logger.debug("mqtt_protocol.handle_publish_packet packet:%s", message) 

132 

133 if not message.qos: 

134 self._context.metrics.on_message_received(0, False) 

135 self._enqueue(message) 

136 return 

137 

138 packet_identifier = cast(int, message.packet_identifier) 

139 state = self._state 

140 

141 # a re-sent QoS 2 message already holds its slot 

142 if packet_identifier not in state.inflight: 

143 if len(state.inflight) >= self._receive_maximum: 

144 raise ReceiveMaximumExceededError( 

145 f"More than {self._receive_maximum} incoming messages in flight" 

146 ) 

147 

148 state.inflight.add(packet_identifier) 

149 

150 if packet_identifier in state.pending_acks: 

151 # the server sent the message again before the ack: the pending 

152 # ack answers both, the message isn't delivered again (MQTT 5, 

153 # 4.3.3) 

154 logger.debug( 

155 "mqtt_protocol.handle_publish_packet.duplicate pid:%s", 

156 packet_identifier, 

157 ) 

158 self._context.metrics.on_message_received(message.qos, True) 

159 return 

160 

161 duplicate = message.qos == 2 and await self._session.has_incoming_message( 

162 packet_identifier 

163 ) 

164 self._context.metrics.on_message_received(message.qos, duplicate) 

165 

166 pending = _PendingAck(message, self._context.now()) 

167 state.pending_acks[packet_identifier] = pending 

168 

169 if duplicate: 

170 # PUBREC was sent already, the message isn't delivered again 

171 logger.debug( 

172 "mqtt_protocol.handle_publish_packet.duplicate pid:%s", 

173 packet_identifier, 

174 ) 

175 pending.acked = True 

176 await self._send_acks() 

177 return 

178 

179 self._enqueue(message) 

180 

181 async def handle_pubrel(self, pubrel_result: PubRelResult) -> None: 

182 connection = self._context.connection 

183 assert connection 

184 

185 logger.debug("mqtt_protocol.handle_pubrel_packet packet:%s", pubrel_result) 

186 

187 packet_identifier = pubrel_result.packet_identifier 

188 

189 self._state.inflight.discard(packet_identifier) 

190 

191 if await self._session.complete_incoming_message(packet_identifier): 

192 reason_code = 0x0 

193 else: 

194 logger.warning( 

195 "mqtt_protocol.handle_pubrel_packet.unknown pid:%s", packet_identifier 

196 ) 

197 reason_code = PACKET_IDENTIFIER_NOT_FOUND 

198 

199 await self._context.write( 

200 connection, pack_pubcomp_packet(packet_identifier, reason_code, {}) 

201 ) 

202 

203 async def ack(self, message: MQTTMessage, reason_code: int = 0) -> None: 

204 """ 

205 Acknowledges a message: PUBACK for QoS 1, PUBREC for QoS 2, nothing for 

206 QoS 0. Acks are sent in the order messages were received (MQTT 5, 4.6), 

207 so an ack waits for the acks of earlier messages: one message which 

208 isn't acknowledged holds back the acks of all later ones, and they 

209 count to "Receive Maximum". A warning is logged when such a message 

210 isn't acknowledged for ACK_BLOCKED_WARNING seconds. 

211 

212 :param reason_code: 0 or >= 0x80 to reject the message; the server 

213 doesn't re-send a rejected message 

214 """ 

215 if reason_code and reason_code < FAILURE_REASON_CODE: 

216 raise ValueError(f"Invalid reason code: 0x{reason_code:02X}") 

217 

218 if not message.qos: 

219 return 

220 

221 pending = self._state.pending_acks.get(cast(int, message.packet_identifier)) 

222 

223 if pending is None or pending.message is not message: 

224 # acknowledged already or received by the previous connection 

225 logger.debug("mqtt_protocol.ack.ignored pid:%s", message.packet_identifier) 

226 return 

227 

228 pending.acked = True 

229 pending.reason_code = reason_code 

230 

231 await self._send_acks() 

232 

233 async def _send_acks(self) -> None: 

234 """Sends acks of the first messages in the order of receiving.""" 

235 async with self._ack_lock: 

236 state = self._state 

237 pending_acks = state.pending_acks 

238 

239 while pending_acks: 

240 packet_identifier, pending = next(iter(pending_acks.items())) 

241 

242 if not pending.acked: 

243 self._ack_monitor.watch( 

244 pending.message, 

245 pending.received_at, 

246 partial(self._messages_waiting_for, state, pending), 

247 ) 

248 return 

249 

250 del pending_acks[packet_identifier] 

251 

252 connection = self._context.connection 

253 

254 if ( 254 ↛ 259line 254 didn't jump to line 259 because the condition on line 254 was never true

255 state is not self._state 

256 or not connection 

257 or connection.is_closing() 

258 ): 

259 return 

260 

261 reason_code = pending.reason_code 

262 

263 if pending.message.qos == 1: 

264 state.inflight.discard(packet_identifier) 

265 packet = pack_puback_packet(packet_identifier, reason_code, {}) 

266 else: 

267 if reason_code < FAILURE_REASON_CODE: 

268 # in flight until PUBREL 

269 await self._session.register_incoming_message(packet_identifier) 

270 else: 

271 # the flow ends with the rejecting PUBREC 

272 state.inflight.discard(packet_identifier) 

273 

274 packet = pack_pubrec_packet(packet_identifier, reason_code, {}) 

275 

276 logger.debug( 

277 "mqtt_protocol.send_ack pid:%s qos:%s reason:0x%02X", 

278 packet_identifier, 

279 pending.message.qos, 

280 reason_code, 

281 ) 

282 await self._context.write(connection, packet) 

283 

284 def _messages_waiting_for( 

285 self, state: _ConnectionState, head: _PendingAck 

286 ) -> Optional[int]: 

287 """ 

288 The number of messages whose acks wait for `head`, None if it doesn't 

289 hold them back anymore. 

290 """ 

291 packet_identifier = cast(int, head.message.packet_identifier) 

292 

293 if ( 

294 state is not self._state 

295 or head.acked 

296 or state.pending_acks.get(packet_identifier) is not head 

297 ): 

298 return None 

299 

300 return len(state.pending_acks) - 1 

301 

302 def _resolve_topic_alias(self, message: MQTTMessage) -> MQTTMessage: 

303 topic_alias = message.properties.get("topic_alias") 

304 

305 if topic_alias is None: 

306 if not message.topic: 

307 raise ProtocolError("PUBLISH without topic and topic alias") 

308 

309 return message 

310 

311 if not 0 < topic_alias <= self._topic_alias_maximum: 

312 raise TopicAliasInvalidError( 

313 f"Topic alias {topic_alias} isn't in range " 

314 f"1..{self._topic_alias_maximum}" 

315 ) 

316 

317 if message.topic: 

318 self._state.topic_aliases[topic_alias] = message.topic 

319 return message 

320 

321 if (topic := self._state.topic_aliases.get(topic_alias)) is None: 

322 raise ProtocolError(f"Unknown topic alias {topic_alias}") 

323 

324 return dataclasses.replace(message, topic=topic) 

325 

326 def _enqueue(self, message: Optional[MQTTMessage]) -> None: 

327 self._buffer.append(message) 

328 self._buffer_ready.set() 

329 self.buffer_monitor.update(len(self._buffer))