Coverage for src/lexigram/graphql/subscriptions/manager.py: 49%

1"""Subscription Manager for GraphQL subscriptions. 

2 

3Manages active subscription streams per connection. 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import contextlib 

10from typing import TYPE_CHECKING, Any 

11 

12from lexigram.logging import get_logger 

13 

14if TYPE_CHECKING: 

15 from collections.abc import AsyncGenerator 

16 

17logger = get_logger(__name__) 

18 

19 

20class SubscriptionManager: 

21 """Manages active subscription streams per connection. 

22 

23 Tracks subscriptions per connection and provides lifecycle management. 

24 Supports optional keepalive pings to prevent load-balancer idle timeouts 

25 on long-lived WebSocket connections. 

26 """ 

27 

28 def __init__(self, keepalive_interval: float = 25.0) -> None: 

29 # connection_id -> {subscription_id -> task} 

30 self._subscriptions: dict[str, dict[str, asyncio.Task]] = {} 

31 # connection_id -> {subscription_id -> AsyncIterator} 

32 self._streams: dict[str, dict[str, AsyncGenerator]] = {} 

33 # connection_id -> keepalive task 

34 self._keepalive_tasks: dict[str, asyncio.Task] = {} 

35 self._keepalive_interval = keepalive_interval 

36 

37 async def subscribe( 

38 self, 

39 connection_id: str, 

40 subscription_id: str, 

41 stream: AsyncGenerator, 

42 ) -> None: 

43 """Register a new subscription. 

44 

45 Args: 

46 connection_id: The WebSocket connection ID. 

47 subscription_id: The unique subscription ID. 

48 stream: The async generator yielding subscription results. 

49 """ 

50 if connection_id not in self._subscriptions: 

51 self._subscriptions[connection_id] = {} 

52 self._streams[connection_id] = {} 

53 

54 # Create task to run the stream 

55 async def run_stream() -> None: 

56 try: 

57 async for _ in stream: 

58 pass # Results are consumed; clients pull from _streams directly 

59 except asyncio.CancelledError: 

60 logger.debug("Subscription %s cancelled", subscription_id) 

61 raise 

62 except Exception as _sub_err: # noqa: BLE001 — subscription stream workers must log any error before re-raising 

63 logger.exception("Error in subscription %s", subscription_id) 

64 raise 

65 

66 task: asyncio.Task[None] = asyncio.create_task(run_stream()) 

67 self._subscriptions[connection_id][subscription_id] = task 

68 self._streams[connection_id][subscription_id] = stream 

69 

70 logger.debug( 

71 "Subscription %s started for connection %s", subscription_id, connection_id 

72 ) 

73 

74 async def unsubscribe( 

75 self, 

76 connection_id: str, 

77 subscription_id: str, 

78 ) -> None: 

79 """Unsubscribe from a specific subscription. 

80 

81 Args: 

82 connection_id: The WebSocket connection ID. 

83 subscription_id: The subscription ID to cancel. 

84 """ 

85 if connection_id not in self._subscriptions: 

86 return 

87 

88 if subscription_id in self._subscriptions[connection_id]: 

89 task = self._subscriptions[connection_id][subscription_id] 

90 task.cancel() 

91 

92 with contextlib.suppress(asyncio.CancelledError): 

93 await task 

94 

95 del self._subscriptions[connection_id][subscription_id] 

96 

97 if subscription_id in self._streams[connection_id]: 

98 del self._streams[connection_id][subscription_id] 

99 

100 logger.debug( 

101 "Subscription %s stopped for connection %s", 

102 subscription_id, 

103 connection_id, 

104 ) 

105 

106 # Clean up empty connections 

107 if not self._subscriptions[connection_id]: 

108 del self._subscriptions[connection_id] 

109 if connection_id in self._streams: 

110 del self._streams[connection_id] 

111 

112 async def disconnect(self, connection_id: str) -> None: 

113 """Clean up all subscriptions and keepalive for a connection. 

114 

115 Args: 

116 connection_id: The WebSocket connection ID to disconnect. 

117 """ 

118 await self.stop_keepalive(connection_id) 

119 

120 if connection_id not in self._subscriptions: 

121 return 

122 

123 # Cancel all subscriptions 

124 for subscription_id in list(self._subscriptions[connection_id].keys()): 

125 await self.unsubscribe(connection_id, subscription_id) 

126 

127 logger.debug("All subscriptions cleaned up for connection %s", connection_id) 

128 

129 async def start_keepalive( 

130 self, 

131 connection_id: str, 

132 send_ping: Any, 

133 ) -> None: 

134 """Start a keepalive loop for a WebSocket connection. 

135 

136 Sends a ping every ``keepalive_interval`` seconds so load-balancer 

137 idle-connection timeouts do not silently drop long-lived subscriptions. 

138 

139 Args: 

140 connection_id: The WebSocket connection ID. 

141 send_ping: Async callable that sends the keepalive message to the 

142 client (e.g., a WebSocket ``send_json`` partial). It will be 

143 called with no arguments on each tick. 

144 """ 

145 if connection_id in self._keepalive_tasks: 

146 return # already running 

147 

148 task: asyncio.Task[None] = asyncio.create_task( 

149 self._keepalive_loop(connection_id, send_ping), 

150 name=f"keepalive:{connection_id}", 

151 ) 

152 self._keepalive_tasks[connection_id] = task 

153 task.add_done_callback( 

154 lambda _t: self._keepalive_tasks.pop(connection_id, None) 

155 ) 

156 logger.debug("Keepalive started for connection %s", connection_id) 

157 

158 async def stop_keepalive(self, connection_id: str) -> None: 

159 """Cancel the keepalive loop for a connection. 

160 

161 Args: 

162 connection_id: The WebSocket connection ID. 

163 """ 

164 task = self._keepalive_tasks.pop(connection_id, None) 

165 if task is not None: 

166 task.cancel() 

167 with contextlib.suppress(asyncio.CancelledError): 

168 await task 

169 logger.debug("Keepalive stopped for connection %s", connection_id) 

170 

171 async def _keepalive_loop(self, connection_id: str, send_ping: Any) -> None: 

172 """Send periodic pings until cancelled. 

173 

174 Args: 

175 connection_id: Used only for log messages. 

176 send_ping: Async callable invoked on each keepalive tick. 

177 """ 

178 while True: 

179 await asyncio.sleep(self._keepalive_interval) 

180 try: 

181 await send_ping() 

182 logger.debug("Keepalive ping sent for connection %s", connection_id) 

183 except asyncio.CancelledError: 

184 raise 

185 except (OSError, RuntimeError) as e: 

186 logger.warning( 

187 "Keepalive ping failed for connection %s; stopping loop: %s", 

188 connection_id, 

189 e, 

190 ) 

191 break 

192 

193 def get_active_subscriptions(self, connection_id: str) -> list[str]: 

194 """Get list of active subscription IDs for a connection. 

195 

196 Args: 

197 connection_id: The WebSocket connection ID. 

198 

199 Returns: 

200 List of active subscription IDs. 

201 """ 

202 if connection_id not in self._subscriptions: 

203 return [] 

204 return list(self._subscriptions[connection_id].keys()) 

205 

206 @property 

207 def active_count(self) -> int: 

208 """Get total number of active subscriptions across all connections.""" 

209 return sum(len(subs) for subs in self._subscriptions.values()) 

210 

211 @property 

212 def connection_count(self) -> int: 

213 """Get number of active connections.""" 

214 return len(self._subscriptions) 

215 

216 

217__all__ = ["SubscriptionManager"]