Coverage for zenmqtt/reconnect.py: 97%

1import asyncio 

2import random 

3from dataclasses import dataclass 

4from logging import getLogger 

5from typing import Awaitable, Callable, Final, Iterator, Optional 

6 

7from zenmqtt.exceptions import ClientConnectionError, ConnectionLostError 

8from zenmqtt.instrumentation.metrics import MetricsCollector 

9from zenmqtt.mqtt.connect import ConnectionResult 

10from zenmqtt.mqtt.protocol import MQTTProtocol 

11from zenmqtt.mqtt.reason_codes import FAILURE_REASON_CODE 

12 

13logger = getLogger(__name__) 

14 

15# CONNACK reason codes which mean that a retry can't help: the client has to 

16# change its configuration (credentials, client id, will, ...) 

17NON_RETRYABLE_CONNACK_REASON_CODES: Final[frozenset[int]] = frozenset( 

18 { 

19 0x81, # Malformed Packet 

20 0x82, # Protocol Error 

21 0x84, # Unsupported Protocol Version 

22 0x85, # Client Identifier not valid 

23 0x86, # Bad User Name or Password 

24 0x87, # Not authorized 

25 0x8A, # Banned 

26 0x8C, # Bad authentication method 

27 0x90, # Topic Name invalid (will topic) 

28 0x95, # Packet too large 

29 0x99, # Payload format invalid (will payload) 

30 0x9A, # Retain not supported (will retain) 

31 0x9B, # QoS not supported (will QoS) 

32 0x9C, # Use another server 

33 0x9D, # Server moved 

34 } 

35) 

36 

37# DISCONNECT reason codes of the server after which the client doesn't reconnect 

38NON_RETRYABLE_DISCONNECT_REASON_CODES: Final[frozenset[int]] = frozenset( 

39 { 

40 0x87, # Not authorized 

41 # another client connected with the same client id: reconnecting would 

42 # take the session back, and both clients would kick each other forever 

43 0x8E, # Session taken over 

44 0x9C, # Use another server 

45 0x9D, # Server moved 

46 } 

47) 

48 

49 

50@dataclass(frozen=True, slots=True) 

51class ReconnectPolicy: 

52 """ 

53 Exponential backoff between reconnect attempts: the first attempt is made 

54 after initial_delay, every next delay is multiplied by multiplier up to 

55 max_delay. Each delay is reduced by a random part up to jitter (0.5 means 

56 up to 50%), so many clients don't reconnect to a restarted server at the 

57 same moment. 

58 

59 Times are in seconds. 

60 """ 

61 

62 initial_delay: float = 1.0 

63 max_delay: float = 60.0 

64 multiplier: float = 2.0 

65 jitter: float = 0.5 

66 # None means the client never gives up 

67 max_attempts: Optional[int] = None 

68 # time for TCP/TLS connect and CONNACK of one attempt 

69 connect_timeout: float = 10.0 

70 

71 def __post_init__(self) -> None: 

72 if self.initial_delay < 0 or self.max_delay < self.initial_delay: 

73 raise ValueError("Expected 0 <= initial_delay <= max_delay") 

74 

75 if self.multiplier < 1: 

76 raise ValueError("Expected multiplier >= 1") 

77 

78 if not 0 <= self.jitter <= 1: 

79 raise ValueError("Expected 0 <= jitter <= 1") 

80 

81 if self.max_attempts is not None and self.max_attempts < 1: 

82 raise ValueError("Expected max_attempts >= 1 or None") 

83 

84 if self.connect_timeout <= 0: 

85 raise ValueError("Expected connect_timeout > 0") 

86 

87 def delays(self) -> Iterator[float]: 

88 """Delays before each attempt.""" 

89 delay = self.initial_delay 

90 attempt = 0 

91 

92 while self.max_attempts is None or attempt < self.max_attempts: 

93 yield delay * (1 - self.jitter * random.random()) 

94 

95 attempt += 1 

96 delay = min(delay * self.multiplier, self.max_delay) 

97 

98 

99DEFAULT_RECONNECT_POLICY: Final[ReconnectPolicy] = ReconnectPolicy() 

100 

101 

102class Reconnector: 

103 """ 

104 Waits until the connection is lost and reconnects with exponential backoff 

105 of the policy. It stops after a non-retryable reason code of the server or 

106 when there are no attempts left. 

107 

108 The owner reacts to the events: 

109 

110 - connect: makes one connection attempt (without Clean Start); 

111 - connection_lost: the connection is lost, before the first attempt; 

112 - reconnected: an attempt succeeded; 

113 - gave_up: the reconnector stopped for good. 

114 """ 

115 

116 def __init__( 

117 self, 

118 policy: ReconnectPolicy, 

119 protocol: MQTTProtocol, 

120 metrics: MetricsCollector, 

121 *, 

122 connect: Callable[[], Awaitable[ConnectionResult]], 

123 connection_lost: Callable[[ConnectionLostError], Awaitable[None]], 

124 reconnected: Callable[[ConnectionResult], Awaitable[None]], 

125 gave_up: Callable[[Exception], Awaitable[None]], 

126 ) -> None: 

127 self._policy = policy 

128 self._protocol = protocol 

129 self._metrics = metrics 

130 

131 self._connect = connect 

132 self._connection_lost = connection_lost 

133 self._reconnected = reconnected 

134 self._gave_up = gave_up 

135 

136 self._task: Optional[asyncio.Task[None]] = None 

137 

138 @property 

139 def running(self) -> bool: 

140 return self._task is not None and not self._task.done() 

141 

142 def start(self) -> None: 

143 self._task = asyncio.create_task(self._run(), name="mqtt-client-reconnect") 

144 

145 async def stop(self) -> None: 

146 if task := self._task: 146 ↛ exitline 146 didn't return from function 'stop' because the condition on line 146 was always true

147 self._task = None 

148 task.cancel() 

149 

150 try: 

151 await task 

152 except asyncio.CancelledError: 

153 pass 

154 

155 async def _run(self) -> None: 

156 try: 

157 await self._reconnect_forever() 

158 except Exception as exc: 

159 logger.error("mqtt_client.reconnect.error", exc_info=exc) 

160 await self._give_up(ConnectionLostError()) 

161 

162 async def _reconnect_forever(self) -> None: 

163 while True: 

164 await self._protocol.wait_closed() 

165 

166 server_disconnect = self._protocol.server_disconnect 

167 lost = ConnectionLostError(server_disconnect) 

168 

169 logger.warning("mqtt_client.connection_lost: %s", lost) 

170 await self._connection_lost(lost) 

171 

172 if ( 

173 server_disconnect 

174 and server_disconnect.reason_code 

175 in NON_RETRYABLE_DISCONNECT_REASON_CODES 

176 ): 

177 logger.error( 

178 "mqtt_client.reconnect.stopped: the server closed the " 

179 "connection with reason 0x%02X", 

180 server_disconnect.reason_code, 

181 ) 

182 await self._give_up(lost) 

183 return 

184 

185 if not await self._reconnect_with_backoff(lost): 

186 return 

187 

188 async def _reconnect_with_backoff(self, lost: ConnectionLostError) -> bool: 

189 """Returns False if the reconnector gave up.""" 

190 for attempt, delay in enumerate(self._policy.delays(), start=1): 

191 self._metrics.on_reconnect_attempt(attempt, delay) 

192 await asyncio.sleep(delay) 

193 

194 logger.info("mqtt_client.reconnect.attempt: %s", attempt) 

195 

196 try: 

197 connack = await asyncio.wait_for( 

198 self._connect(), self._policy.connect_timeout 

199 ) 

200 except (OSError, asyncio.TimeoutError, ClientConnectionError) as exc: 

201 logger.warning("mqtt_client.reconnect.failed: %r", exc) 

202 continue 

203 

204 if connack.reason_code >= FAILURE_REASON_CODE: 

205 if connack.reason_code in NON_RETRYABLE_CONNACK_REASON_CODES: 

206 logger.error( 

207 "mqtt_client.reconnect.stopped: the server rejected the " 

208 "connection with reason 0x%02X", 

209 connack.reason_code, 

210 ) 

211 await self._give_up(lost) 

212 return False 

213 

214 logger.warning( 

215 "mqtt_client.reconnect.rejected: reason 0x%02X", 

216 connack.reason_code, 

217 ) 

218 continue 

219 

220 logger.info("mqtt_client.reconnect.succeeded: attempt %s", attempt) 

221 

222 await self._reconnected(connack) 

223 return True 

224 

225 logger.error("mqtt_client.reconnect.stopped: no attempts left") 

226 await self._give_up(lost) 

227 return False 

228 

229 async def _give_up(self, exc: Exception) -> None: 

230 # the task ends after this, the owner may connect again 

231 self._task = None 

232 self._metrics.on_reconnect_gave_up() 

233 

234 await self._gave_up(exc)