Rhinox13/chatapi
0
1from __future__ import annotations
2
3from dataclasses import dataclass, field
4import queue
5import threading
6import uuid
7from typing import Any, Callable
8
9from ..repositories import ConversationStore, UserStore
10
11
12@dataclass
13class RealtimeSubscription:
14 owner_id: str
15 events: queue.Queue[dict[str, Any]] = field(default_factory=lambda: queue.Queue(maxsize=100))
16 closed: threading.Event = field(default_factory=threading.Event)
17
18
19@dataclass
20class ActiveConnection:
21 connection_id: str
22 owner_id: str
23 kind: str
24 close_callback: Callable[[], None] | None = None
25
26
27@dataclass(frozen=True)
28class ConnectionLease:
29 connection_id: str
30 owner_id: str
31 kind: str
32
33
34class ConnectionLimitExceeded(RuntimeError):
35 pass
36
37
38class RealtimeBroker:
39 def __init__(self, store: ConversationStore, user_store: UserStore):
40 self._store = store
41 self._user_store = user_store
42 self._lock = threading.Lock()
43 self._subscriptions_by_owner: dict[str, list[RealtimeSubscription]] = {}
44 self._connections_by_owner: dict[str, list[ActiveConnection]] = {}
45
46 @staticmethod
47 def _normalize_limit(value: Any, default: int = 0) -> int:
48 try:
49 return max(0, int(value or default))
50 except (TypeError, ValueError):
51 return default
52
53 def acquire_connection(
54 self,
55 owner_id: str,
56 *,
57 kind: str,
58 max_connections: int = 0,
59 max_connections_per_user: int = 0,
60 close_callback: Callable[[], None] | None = None,
61 ) -> ConnectionLease:
62 connection = ActiveConnection(
63 connection_id=uuid.uuid4().hex,
64 owner_id=owner_id,
65 kind=kind,
66 close_callback=close_callback,
67 )
68 affected_owners: set[str] = set()
69 with self._lock:
70 affected_owners |= self._enforce_limits_locked(
71 owner_id,
72 max_connections=max_connections,
73 max_connections_per_user=max_connections_per_user,
74 )
75 if self._would_exceed_limits_locked(
76 owner_id,
77 max_connections=max_connections,
78 max_connections_per_user=max_connections_per_user,
79 ):
80 raise ConnectionLimitExceeded("connection limit exceeded")
81 self._connections_by_owner.setdefault(owner_id, []).append(connection)
82 affected_owners.add(owner_id)
83 self._publish_connection_counts(affected_owners)
84 return ConnectionLease(
85 connection_id=connection.connection_id,
86 owner_id=owner_id,
87 kind=kind,
88 )
89
90 def subscribe(
91 self,
92 owner_id: str,
93 *,
94 max_connections: int = 0,
95 max_connections_per_user: int = 0,
96 queue_size: int = 100,
97 ) -> RealtimeSubscription:
98 queue_size = max(1, self._normalize_limit(queue_size, 100))
99 subscription = RealtimeSubscription(
100 owner_id=owner_id,
101 events=queue.Queue(maxsize=queue_size),
102 )
103 lease = self.acquire_connection(
104 owner_id,
105 kind="websocket",
106 max_connections=max_connections,
107 max_connections_per_user=max_connections_per_user,
108 close_callback=subscription.closed.set,
109 )
110 setattr(subscription, "_connection_lease", lease)
111 with self._lock:
112 subscribers = self._subscriptions_by_owner.setdefault(owner_id, [])
113 subscribers.append(subscription)
114 return subscription
115
116 def _enforce_limits_locked(
117 self,
118 owner_id: str,
119 *,
120 max_connections: int,
121 max_connections_per_user: int,
122 ) -> set[str]:
123 affected_owners: set[str] = set()
124 max_connections = self._normalize_limit(max_connections)
125 max_connections_per_user = self._normalize_limit(max_connections_per_user)
126
127 if max_connections_per_user > 0:
128 while self._owner_connection_count_locked(owner_id) >= max_connections_per_user:
129 removed = self._drop_oldest_connection_locked(owner_id=owner_id)
130 if removed is None:
131 break
132 affected_owners.add(removed.owner_id)
133
134 if max_connections <= 0:
135 return affected_owners
136 while self._connection_count_locked() >= max_connections:
137 removed = self._drop_oldest_connection_locked()
138 if removed is None:
139 return affected_owners
140 affected_owners.add(removed.owner_id)
141 return affected_owners
142
143 def _connection_count_locked(self) -> int:
144 return sum(len(items) for items in self._connections_by_owner.values())
145
146 def _owner_connection_count_locked(self, owner_id: str) -> int:
147 return len(self._connections_by_owner.get(owner_id, ()))
148
149 def _would_exceed_limits_locked(
150 self,
151 owner_id: str,
152 *,
153 max_connections: int,
154 max_connections_per_user: int,
155 ) -> bool:
156 max_connections = self._normalize_limit(max_connections)
157 max_connections_per_user = self._normalize_limit(max_connections_per_user)
158 if max_connections_per_user > 0 and self._owner_connection_count_locked(owner_id) >= max_connections_per_user:
159 return True
160 if max_connections > 0 and self._connection_count_locked() >= max_connections:
161 return True
162 return False
163
164 def _drop_oldest_connection_locked(self, owner_id: str | None = None) -> ActiveConnection | None:
165 candidate_owners = [owner_id] if owner_id is not None else list(self._connections_by_owner.keys())
166 for candidate_owner in candidate_owners:
167 connections = self._connections_by_owner.get(candidate_owner, [])
168 for index, connection in enumerate(connections):
169 if connection.close_callback is None:
170 continue
171 connections.pop(index)
172 if connections:
173 self._connections_by_owner[candidate_owner] = connections
174 else:
175 self._connections_by_owner.pop(candidate_owner, None)
176 self._close_connection_locked(connection)
177 return connection
178 return None
179
180 def _close_connection_locked(self, connection: ActiveConnection) -> None:
181 self._remove_subscription_for_connection_locked(connection.connection_id)
182 if connection.close_callback is not None:
183 connection.close_callback()
184
185 def _remove_subscription_for_connection_locked(self, connection_id: str) -> None:
186 for owner_id, subscribers in list(self._subscriptions_by_owner.items()):
187 remaining = [
188 subscription
189 for subscription in subscribers
190 if getattr(getattr(subscription, "_connection_lease", None), "connection_id", None) != connection_id
191 ]
192 if len(remaining) == len(subscribers):
193 continue
194 if remaining:
195 self._subscriptions_by_owner[owner_id] = remaining
196 else:
197 self._subscriptions_by_owner.pop(owner_id, None)
198 return
199
200 def unsubscribe(self, subscription: RealtimeSubscription) -> None:
201 lease = getattr(subscription, "_connection_lease", None)
202 if isinstance(lease, ConnectionLease):
203 self.release_connection(lease)
204 return
205 with self._lock:
206 subscribers = self._subscriptions_by_owner.get(subscription.owner_id)
207 if not subscribers:
208 return
209 self._subscriptions_by_owner[subscription.owner_id] = [
210 item for item in subscribers if item is not subscription
211 ]
212 if not self._subscriptions_by_owner[subscription.owner_id]:
213 self._subscriptions_by_owner.pop(subscription.owner_id, None)
214
215 def release_connection(self, lease: ConnectionLease) -> None:
216 affected_owner: str | None = None
217 with self._lock:
218 connections = self._connections_by_owner.get(lease.owner_id)
219 if not connections:
220 return
221 remaining = [item for item in connections if item.connection_id != lease.connection_id]
222 if len(remaining) == len(connections):
223 return
224 affected_owner = lease.owner_id
225 if remaining:
226 self._connections_by_owner[lease.owner_id] = remaining
227 else:
228 self._connections_by_owner.pop(lease.owner_id, None)
229 self._remove_subscription_for_connection_locked(lease.connection_id)
230 if affected_owner is not None:
231 self._publish_connection_counts({affected_owner})
232
233 def count_owner_connections(self, owner_id: str) -> int:
234 with self._lock:
235 return self._owner_connection_count_locked(owner_id)
236
237 def count_connections(self) -> int:
238 with self._lock:
239 return self._connection_count_locked()
240
241 def publish_snapshot(self, owner_id: str) -> None:
242 self._publish(owner_id, self.build_snapshot(owner_id))
243
244 def publish_conversation_upsert(self, owner_id: str, conversation_id: str) -> None:
245 conversation = self._store.get_conversation(conversation_id, owner_id)
246 if conversation is None:
247 self.publish_conversation_delete(owner_id, conversation_id)
248 return
249 try:
250 messages = self._store.get_messages(conversation_id, owner_id)
251 except ValueError:
252 messages = []
253 self._publish(
254 owner_id,
255 {
256 "type": "conversation_upsert",
257 "conversation": conversation.to_dict(),
258 "messages": [message.to_dict() for message in messages],
259 },
260 )
261
262 def publish_conversation_delete(self, owner_id: str, conversation_id: str) -> None:
263 self._publish(
264 owner_id,
265 {
266 "type": "conversation_delete",
267 "conversation_id": conversation_id,
268 },
269 )
270
271 def _publish(self, owner_id: str, event: dict[str, Any]) -> None:
272 with self._lock:
273 subscribers = tuple(self._subscriptions_by_owner.get(owner_id, ()))
274 for subscription in subscribers:
275 self._offer_event(subscription, event)
276
277 def _publish_connection_counts(self, owner_ids: set[str]) -> None:
278 for owner_id in owner_ids:
279 self._publish(
280 owner_id,
281 {
282 "type": "connection_count",
283 "current_connection_count": self.count_owner_connections(owner_id),
284 },
285 )
286
287 @staticmethod
288 def _offer_event(subscription: RealtimeSubscription, event: dict[str, Any]) -> None:
289 try:
290 subscription.events.put_nowait(event)
291 return
292 except queue.Full:
293 pass
294 try:
295 subscription.events.get_nowait()
296 except queue.Empty:
297 pass
298 try:
299 subscription.events.put_nowait(event)
300 except queue.Full:
301 subscription.closed.set()
302
303 @staticmethod
304 def _force_event(subscription: RealtimeSubscription, event: dict[str, Any]) -> None:
305 while True:
306 try:
307 subscription.events.put_nowait(event)
308 return
309 except queue.Full:
310 try:
311 subscription.events.get_nowait()
312 except queue.Empty:
313 return
314
315 def build_snapshot(
316 self,
317 owner_id: str,
318 ) -> dict[str, Any]:
319 return {
320 "type": "snapshot",
321 "conversations": [item.to_dict() for item in self._store.list_conversations(owner_id)],
322 }
323