Coverage for zenmqtt/session.py: 100%

1import abc 

2import asyncio 

3import dataclasses 

4import logging 

5from dataclasses import dataclass 

6from enum import IntEnum 

7from typing import Final, Optional, Protocol, Sequence, TypedDict 

8 

9from zenmqtt.mqtt.publish import PublishProperties 

10 

11PacketIdentifier = int 

12 

13MAX_PACKET_IDENTIFIER: Final[int] = 2**16 - 1 

14 

15logger = logging.getLogger(__name__) 

16 

17 

18class OutgoingMessageState(IntEnum): 

19 # PUBLISH was sent, waiting for PUBACK (QoS 1) or PUBREC (QoS 2) 

20 AWAITING_ACK = 1 

21 # PUBREL was sent, waiting for PUBCOMP (QoS 2) 

22 AWAITING_COMP = 2 

23 

24 

25@dataclass(frozen=True) 

26class OutgoingMessage: 

27 __slots__ = ( 

28 "packet_identifier", 

29 "topic", 

30 "payload", 

31 "qos", 

32 "retain", 

33 "properties", 

34 "state", 

35 ) 

36 

37 packet_identifier: PacketIdentifier 

38 topic: str 

39 payload: bytes 

40 qos: int 

41 retain: bool 

42 properties: PublishProperties 

43 state: OutgoingMessageState 

44 

45 

46class MQTTSession(Protocol): 

47 """ 

48 Client side session state (MQTT 5, section 4.1): 

49 

50 * QoS 1 and QoS 2 messages which have been sent to the server, 

51 but have not been completely acknowledged; 

52 * QoS 2 messages which have been received from the server, 

53 but have not been completely acknowledged. 

54 

55 The session also owns packet identifiers: an identifier stays acquired 

56 while its outgoing message is pending, so a persistent implementation 

57 must not hand out identifiers of restored pending messages. 

58 

59 Custom storages should rather inherit from BaseSession. 

60 """ 

61 

62 async def reset(self) -> None: 

63 """Discard the whole session state.""" 

64 ... 

65 

66 async def acquire_packet_identifier(self) -> PacketIdentifier: 

67 """A free packet identifier; waits while all of them are in use.""" 

68 ... 

69 

70 async def release_packet_identifier( 

71 self, packet_identifier: PacketIdentifier 

72 ) -> None: 

73 """Release identifier which isn't bound to a message (SUBSCRIBE, ...).""" 

74 ... 

75 

76 async def store_outgoing_message(self, message: OutgoingMessage) -> None: 

77 """Called before the PUBLISH packet with QoS > 0 is sent.""" 

78 ... 

79 

80 async def mark_outgoing_message_released( 

81 self, packet_identifier: PacketIdentifier 

82 ) -> bool: 

83 """ 

84 Called on successful PUBREC; the message keeps its original order. 

85 Returns False if there is no such message. 

86 """ 

87 ... 

88 

89 async def complete_outgoing_message( 

90 self, packet_identifier: PacketIdentifier 

91 ) -> None: 

92 """ 

93 Called on PUBACK, PUBCOMP or PUBREC with an error reason code. 

94 Removes the message and releases its packet identifier. 

95 """ 

96 ... 

97 

98 async def get_pending_outgoing_messages(self) -> Sequence[OutgoingMessage]: 

99 """Pending messages in the order they were stored.""" 

100 ... 

101 

102 async def register_incoming_message( 

103 self, packet_identifier: PacketIdentifier 

104 ) -> bool: 

105 """ 

106 Called on incoming PUBLISH with QoS 2 when PUBREC is sent. 

107 Returns False if the message was already received (duplicate). 

108 """ 

109 ... 

110 

111 async def has_incoming_message(self, packet_identifier: PacketIdentifier) -> bool: 

112 """ 

113 True if PUBREC was sent for the incoming QoS 2 message and PUBREL 

114 isn't received yet. 

115 """ 

116 ... 

117 

118 async def complete_incoming_message( 

119 self, packet_identifier: PacketIdentifier 

120 ) -> bool: 

121 """ 

122 Called on incoming PUBREL. 

123 Returns False if the message wasn't registered. 

124 """ 

125 ... 

126 

127 

128class SessionOptions(TypedDict, total=False): 

129 # how long to wait for a free packet identifier, seconds 

130 timeout: float 

131 

132 

133_DEFAULT_SESSION_OPTIONS: Final[SessionOptions] = { 

134 "timeout": 10, 

135} 

136 

137 

138class BaseSession(MQTTSession, abc.ABC): 

139 """ 

140 Common session logic: packet identifiers, state transitions, duplicates 

141 detection. Implementations only provide the storage primitives below. 

142 

143 The packet identifiers pool is built on the first use and skips 

144 identifiers of the stored pending messages, so a restored session 

145 doesn't reuse them. 

146 """ 

147 

148 def __init__(self, options: Optional[SessionOptions] = None) -> None: 

149 self._options: SessionOptions = { 

150 **_DEFAULT_SESSION_OPTIONS, 

151 **(options or {}), 

152 } 

153 

154 self._packet_identifiers_pool: Optional[asyncio.Queue[PacketIdentifier]] = None 

155 self._packet_identifiers_pool_lock = asyncio.Lock() 

156 self._acquired_packet_identifiers: set[PacketIdentifier] = set() 

157 

158 async def reset(self) -> None: 

159 logger.debug("mqtt_session.reset") 

160 

161 await self._clear() 

162 

163 # return identifiers to the same pool, so pending acquirers are not lost 

164 for packet_identifier in list(self._acquired_packet_identifiers): 

165 self._release(packet_identifier) 

166 

167 async def acquire_packet_identifier(self) -> PacketIdentifier: 

168 pool = await self._get_packet_identifiers_pool() 

169 

170 packet_identifier = await asyncio.wait_for( 

171 pool.get(), timeout=self._options["timeout"] 

172 ) 

173 self._acquired_packet_identifiers.add(packet_identifier) 

174 

175 logger.debug("mqtt_session.acquire_packet_identifier pid:%s", packet_identifier) 

176 return packet_identifier 

177 

178 async def release_packet_identifier( 

179 self, packet_identifier: PacketIdentifier 

180 ) -> None: 

181 logger.debug("mqtt_session.release_packet_identifier pid:%s", packet_identifier) 

182 self._release(packet_identifier) 

183 

184 async def store_outgoing_message(self, message: OutgoingMessage) -> None: 

185 logger.debug( 

186 "mqtt_session.store_outgoing_message pid:%s", message.packet_identifier 

187 ) 

188 await self._save_outgoing_message(message) 

189 

190 async def mark_outgoing_message_released( 

191 self, packet_identifier: PacketIdentifier 

192 ) -> bool: 

193 logger.debug( 

194 "mqtt_session.mark_outgoing_message_released pid:%s", packet_identifier 

195 ) 

196 

197 if not (message := await self._load_outgoing_message(packet_identifier)): 

198 return False 

199 

200 await self._save_outgoing_message( 

201 dataclasses.replace(message, state=OutgoingMessageState.AWAITING_COMP) 

202 ) 

203 return True 

204 

205 async def complete_outgoing_message( 

206 self, packet_identifier: PacketIdentifier 

207 ) -> None: 

208 logger.debug("mqtt_session.complete_outgoing_message pid:%s", packet_identifier) 

209 

210 if await self._delete_outgoing_message(packet_identifier): 

211 self._release(packet_identifier) 

212 

213 async def get_pending_outgoing_messages(self) -> Sequence[OutgoingMessage]: 

214 return await self._load_outgoing_messages() 

215 

216 async def register_incoming_message( 

217 self, packet_identifier: PacketIdentifier 

218 ) -> bool: 

219 logger.debug("mqtt_session.register_incoming_message pid:%s", packet_identifier) 

220 

221 if await self._has_incoming_message(packet_identifier): 

222 return False 

223 

224 await self._save_incoming_message(packet_identifier) 

225 return True 

226 

227 async def has_incoming_message(self, packet_identifier: PacketIdentifier) -> bool: 

228 return await self._has_incoming_message(packet_identifier) 

229 

230 async def complete_incoming_message( 

231 self, packet_identifier: PacketIdentifier 

232 ) -> bool: 

233 logger.debug("mqtt_session.complete_incoming_message pid:%s", packet_identifier) 

234 

235 if not await self._has_incoming_message(packet_identifier): 

236 return False 

237 

238 await self._delete_incoming_message(packet_identifier) 

239 return True 

240 

241 # storage primitives, implemented by subclasses 

242 

243 @abc.abstractmethod 

244 async def _save_outgoing_message(self, message: OutgoingMessage) -> None: 

245 """Insert the message or update it in place, keeping its order.""" 

246 

247 @abc.abstractmethod 

248 async def _load_outgoing_message( 

249 self, packet_identifier: PacketIdentifier 

250 ) -> Optional[OutgoingMessage]: 

251 """None if there is no such message.""" 

252 

253 @abc.abstractmethod 

254 async def _load_outgoing_messages(self) -> Sequence[OutgoingMessage]: 

255 """All stored messages in the order they were inserted.""" 

256 

257 @abc.abstractmethod 

258 async def _delete_outgoing_message( 

259 self, packet_identifier: PacketIdentifier 

260 ) -> bool: 

261 """Returns False if there was no such message.""" 

262 

263 @abc.abstractmethod 

264 async def _save_incoming_message(self, packet_identifier: PacketIdentifier) -> None: 

265 """Saves the identifier of an incoming QoS 2 message.""" 

266 

267 @abc.abstractmethod 

268 async def _has_incoming_message(self, packet_identifier: PacketIdentifier) -> bool: 

269 """True if the identifier is saved.""" 

270 

271 @abc.abstractmethod 

272 async def _delete_incoming_message( 

273 self, packet_identifier: PacketIdentifier 

274 ) -> None: 

275 """Deletes the identifier of an incoming QoS 2 message.""" 

276 

277 @abc.abstractmethod 

278 async def _clear(self) -> None: 

279 """Delete all outgoing and incoming messages.""" 

280 

281 async def _get_packet_identifiers_pool(self) -> asyncio.Queue[PacketIdentifier]: 

282 async with self._packet_identifiers_pool_lock: 

283 if self._packet_identifiers_pool is None: 

284 pending = { 

285 message.packet_identifier 

286 for message in await self._load_outgoing_messages() 

287 } 

288 

289 pool: asyncio.Queue[PacketIdentifier] = asyncio.Queue( 

290 maxsize=MAX_PACKET_IDENTIFIER 

291 ) 

292 

293 for packet_identifier in range(1, MAX_PACKET_IDENTIFIER + 1): 

294 if packet_identifier not in pending: 

295 pool.put_nowait(packet_identifier) 

296 

297 self._acquired_packet_identifiers.update(pending) 

298 self._packet_identifiers_pool = pool 

299 

300 return self._packet_identifiers_pool 

301 

302 def _release(self, packet_identifier: PacketIdentifier) -> None: 

303 # releasing is idempotent, so the pool never contains duplicates 

304 if packet_identifier in self._acquired_packet_identifiers: 

305 assert self._packet_identifiers_pool 

306 

307 self._acquired_packet_identifiers.remove(packet_identifier) 

308 self._packet_identifiers_pool.put_nowait(packet_identifier) 

309 

310 

311class InMemorySession(BaseSession): 

312 def __init__(self, options: Optional[SessionOptions] = None) -> None: 

313 super().__init__(options) 

314 

315 # dict keeps the insertion order, re-assigning a key doesn't change it 

316 self._outgoing_messages: dict[PacketIdentifier, OutgoingMessage] = {} 

317 self._incoming_messages: set[PacketIdentifier] = set() 

318 

319 async def _save_outgoing_message(self, message: OutgoingMessage) -> None: 

320 self._outgoing_messages[message.packet_identifier] = message 

321 

322 async def _load_outgoing_message( 

323 self, packet_identifier: PacketIdentifier 

324 ) -> Optional[OutgoingMessage]: 

325 return self._outgoing_messages.get(packet_identifier) 

326 

327 async def _load_outgoing_messages(self) -> Sequence[OutgoingMessage]: 

328 return list(self._outgoing_messages.values()) 

329 

330 async def _delete_outgoing_message( 

331 self, packet_identifier: PacketIdentifier 

332 ) -> bool: 

333 return self._outgoing_messages.pop(packet_identifier, None) is not None 

334 

335 async def _save_incoming_message(self, packet_identifier: PacketIdentifier) -> None: 

336 self._incoming_messages.add(packet_identifier) 

337 

338 async def _has_incoming_message(self, packet_identifier: PacketIdentifier) -> bool: 

339 return packet_identifier in self._incoming_messages 

340 

341 async def _delete_incoming_message( 

342 self, packet_identifier: PacketIdentifier 

343 ) -> None: 

344 self._incoming_messages.discard(packet_identifier) 

345 

346 async def _clear(self) -> None: 

347 self._outgoing_messages = {} 

348 self._incoming_messages = set() 

349 

350 

351def build_default_session() -> MQTTSession: 

352 return InMemorySession()