Coverage for src/lexigram/web/sse/heartbeat.py: 30%

1"""Shared SSE heartbeat scheduler. 

2 

3Instead of one asyncio task per SSE connection, a single ``SSEHeartbeatScheduler`` 

4task fires heartbeats to all active connections. This keeps the number of live 

5asyncio tasks proportional to the number of **distinct heartbeat intervals** used, 

6not the number of open connections. 

7 

8Usage:: 

9 

10 from lexigram.web.sse.heartbeat import get_heartbeat_scheduler 

11 

12 scheduler = get_heartbeat_scheduler(interval=30.0) 

13 await scheduler.start() # once, at app startup (or lazily) 

14 

15 token = scheduler.register(handler) 

16 # ... stream events ... 

17 scheduler.unregister(token) 

18 

19The ``SSEResponse`` class integrates this automatically. 

20""" 

21 

22from __future__ import annotations 

23 

24import asyncio 

25from typing import Any 

26 

27from lexigram.logging import get_logger 

28from lexigram.primitives import clock as ambient_clock 

29 

30logger = get_logger(__name__) 

31 

32# Heartbeat payload — a comment line keeps the connection alive without 

33# triggering the client's `message` event listener. 

34_HEARTBEAT_PAYLOAD = ": heartbeat\n\n" 

35 

36 

37class SSEHeartbeatScheduler: 

38 """Single-task heartbeat scheduler for multiple SSE connections. 

39 

40 Fires a heartbeat comment to all registered ``SSEBackpressureHandler`` 

41 instances at the configured ``interval``. Uses one background asyncio 

42 task regardless of how many connections are active. 

43 

44 Args: 

45 interval: Seconds between heartbeat messages (default: 30.0). 

46 tick_resolution: How often (in seconds) the scheduler checks whether 

47 heartbeats are due. Lower values give more accurate timing at 

48 the cost of slightly more CPU work (default: 1.0). 

49 """ 

50 

51 def __init__( 

52 self, 

53 interval: float = 30.0, 

54 tick_resolution: float = 1.0, 

55 ) -> None: 

56 self.interval = interval 

57 self.tick_resolution = tick_resolution 

58 # Use a plain set with lock rather than WeakSet so we can track 

59 # last-sent times without keeping the handler alive. 

60 self._handlers: set[Any] = set() 

61 self._last_sent: dict[int, float] = {} # id(handler) → last send time 

62 self._lock = asyncio.Lock() 

63 self._task: asyncio.Task[None] | None = None 

64 self._running = False 

65 

66 async def start(self) -> None: 

67 """Start the background heartbeat loop (idempotent).""" 

68 if self._running: 

69 return 

70 self._running = True 

71 self._task = asyncio.create_task(self._loop(), name="sse-heartbeat-scheduler") 

72 

73 async def stop(self) -> None: 

74 """Stop the background heartbeat loop.""" 

75 self._running = False 

76 if self._task is not None: 

77 self._task.cancel() 

78 try: 

79 await self._task 

80 except asyncio.CancelledError: 

81 pass 

82 self._task = None 

83 

84 async def register(self, handler: Any) -> None: 

85 """Register an SSE connection handler for heartbeats. 

86 

87 Args: 

88 handler: An ``SSEBackpressureHandler`` instance. 

89 """ 

90 async with self._lock: 

91 self._handlers.add(handler) 

92 self._last_sent[id(handler)] = ambient_clock.monotonic() 

93 # Start lazily if not already running 

94 if not self._running: 

95 await self.start() 

96 

97 async def unregister(self, handler: Any) -> None: 

98 """Unregister a handler when the connection closes. 

99 

100 Args: 

101 handler: The previously registered handler. 

102 """ 

103 async with self._lock: 

104 self._handlers.discard(handler) 

105 self._last_sent.pop(id(handler), None) 

106 

107 async def _loop(self) -> None: 

108 """Main heartbeat loop — runs until stopped.""" 

109 while self._running: 

110 await asyncio.sleep(self.tick_resolution) 

111 now = ambient_clock.monotonic() 

112 async with self._lock: 

113 handlers_snapshot = list(self._handlers) 

114 

115 for handler in handlers_snapshot: 

116 hid = id(handler) 

117 last = self._last_sent.get(hid, now) 

118 if now - last >= self.interval: 

119 try: 

120 await handler.send("heartbeat", _HEARTBEAT_PAYLOAD) 

121 async with self._lock: 

122 self._last_sent[hid] = now 

123 except (OSError, RuntimeError): 

124 # Handler may already be closed; ignore silently. 

125 pass 

126 

127 @property 

128 def active_connections(self) -> int: 

129 """Number of SSE connections currently registered.""" 

130 return len(self._handlers) 

131 

132 

133# Module-level cache: interval → scheduler 

134_schedulers: dict[float, SSEHeartbeatScheduler] = {} 

135_scheduler_lock: asyncio.Lock | None = None 

136 

137 

138def get_heartbeat_scheduler( 

139 interval: float = 30.0, 

140) -> SSEHeartbeatScheduler: 

141 """Return (creating if necessary) a shared scheduler for ``interval``. 

142 

143 All SSE connections with the same heartbeat interval share one scheduler 

144 instance and therefore one background asyncio task. 

145 

146 Args: 

147 interval: Heartbeat interval in seconds. 

148 

149 Returns: 

150 The shared ``SSEHeartbeatScheduler`` for this interval. 

151 """ 

152 key = interval 

153 if key not in _schedulers: 

154 _schedulers[key] = SSEHeartbeatScheduler(interval=interval) 

155 return _schedulers[key] 

156 

157 

158__all__ = ["SSEHeartbeatScheduler", "get_heartbeat_scheduler"]