Coverage for zenmqtt/mqtt/protocol/keepalive.py: 95%

1import asyncio 

2from asyncio import Task 

3from logging import getLogger 

4from typing import Optional 

5 

6from zenmqtt.exceptions import ConnectionLostError, NotConnectedError 

7from zenmqtt.mqtt.ping import PingRespResult, pack_pingreq_packet 

8from zenmqtt.mqtt.protocol.context import ProtocolContext 

9 

10logger = getLogger(__name__) 

11 

12 

13class KeepAlive: 

14 """PINGREQ/PINGRESP and keep alive (MQTT 5, 3.1.2.10).""" 

15 

16 def __init__(self, context: ProtocolContext) -> None: 

17 self._context = context 

18 

19 # PINGRESP has no packet identifier, so only one PINGREQ is in flight 

20 self.ping_future: Optional[asyncio.Future[None]] = None 

21 self._ping_sent_at = 0.0 

22 

23 self.task: Optional[Task[None]] = None 

24 

25 async def ping(self) -> None: 

26 connection = self._context.ensure_connected() 

27 

28 # concurrent callers share the in-flight PINGREQ 

29 if (future := self.ping_future) is None: 

30 future = asyncio.get_running_loop().create_future() 

31 self.ping_future = future 

32 

33 logger.debug("mqtt_protocol.send_pingreq_packet") 

34 self._ping_sent_at = self._context.now() 

35 try: 

36 await self._context.write(connection, pack_pingreq_packet()) 

37 except BaseException: 

38 self.ping_future = None 

39 raise 

40 

41 # cancellation of one caller must not affect the others 

42 await asyncio.shield(future) 

43 

44 async def handle_pingresp(self, packet: PingRespResult) -> None: 

45 logger.debug("mqtt_protocol.handle_pingresp_packet") 

46 

47 future, self.ping_future = self.ping_future, None 

48 

49 if future and not future.done(): 

50 self._context.metrics.on_ping(self._context.now() - self._ping_sent_at) 

51 future.set_result(None) 

52 

53 def start(self, keepalive: int) -> None: 

54 """ 

55 :param keepalive: seconds between control packets sent by the client, 

56 0 disables keep alive 

57 """ 

58 if keepalive: 

59 self.task = asyncio.create_task( 

60 self._keep_alive_loop(keepalive), name="mqtt-protocol-keep-alive" 

61 ) 

62 

63 def on_connection_lost(self, exc: Exception) -> None: 

64 if self.task: 

65 self.task.cancel() 

66 self.task = None 

67 

68 future, self.ping_future = self.ping_future, None 

69 

70 if future and not future.done(): 

71 future.set_exception(exc) 

72 

73 async def _keep_alive_loop(self, keepalive: int) -> None: 

74 """ 

75 Sends PINGREQ if the client didn't send any packet within the keep alive 

76 period, and closes the connection if PINGRESP doesn't come in time. 

77 """ 

78 connection = self._context.ensure_connected() 

79 loop = asyncio.get_running_loop() 

80 

81 while True: 

82 last_write_at = connection.last_write_at or loop.time() 

83 delay = last_write_at + keepalive - loop.time() 

84 

85 if delay > 0: 

86 await asyncio.sleep(delay) 

87 continue 

88 

89 try: 

90 await asyncio.wait_for(self.ping(), keepalive) 

91 except asyncio.TimeoutError: 

92 logger.warning("mqtt_protocol.keep_alive.pingresp_timeout") 

93 self._context.metrics.on_ping_timeout() 

94 # the read loop handles the lost connection 

95 await connection.disconnect() 

96 return 

97 except (NotConnectedError, ConnectionLostError): 

98 return