Coverage for /usr/local/lib/python3.10/site-packages/opal_server-0.0.0-py3.10.egg/opal_server/statistics.py: 62%
190 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 11:54 +0000
1import asyncio
2import os
3from datetime import datetime
4from importlib.metadata import version as module_version
5from random import uniform
6from typing import Any, Dict, List, Optional, Set
7from uuid import uuid4
9import opal_server
10import pydantic
11from fastapi import APIRouter, HTTPException, status
12from fastapi_websocket_pubsub.event_notifier import Subscription, TopicList
13from fastapi_websocket_pubsub.pub_sub_server import PubSubEndpoint
14from opal_common.async_utils import TasksPool
15from opal_common.config import opal_common_config
16from opal_common.logger import get_logger
17from opal_common.topics.publisher import PeriodicPublisher
18from opal_server.config import opal_server_config
19from pydantic import BaseModel, Field
22class ChannelStats(BaseModel):
23 rpc_id: str
24 client_id: str
25 topics: TopicList
28class ServerStats(BaseModel):
29 uptime: datetime = Field(..., description="uptime for this opal server worker")
30 version: str = Field(..., description="opal server version")
31 clients: Dict[str, List[ChannelStats]] = Field(
32 ...,
33 description="connected opal clients, each client can have multiple subscriptions",
34 )
35 servers: Set[str] = Field(
36 ...,
37 description="list of all connected opal server replicas",
38 )
41class ServerStatsBrief(BaseModel):
42 uptime: datetime = Field(..., description="uptime for this opal server worker")
43 version: str = Field(..., description="opal server version")
44 client_count: int = Field(..., description="number of connected opal clients")
45 server_count: int = Field(..., description="number of opal server replicas")
48class SyncRequest(BaseModel):
49 requesting_worker_id: str
52class SyncResponse(BaseModel):
53 requesting_worker_id: str
54 clients: Dict[str, List[ChannelStats]]
55 rpc_id_to_client_id: Dict[str, str]
58class ServerKeepalive(BaseModel):
59 worker_id: str
62logger = get_logger("opal.statistics")
64# time to wait before sending statistics
65MIN_TIME_TO_WAIT = 0.001
66MAX_TIME_TO_WAIT = 5
67SLEEP_TIME_FOR_BROADCASTER_READER_TO_START = 2
70class OpalStatistics:
71 """Manage opal server statistics.
73 Args:
74 endpoint:
75 The pub/sub server endpoint that allows us to subscribe to the stats channel on the server side
76 """
78 def __init__(self, endpoint):
79 self._endpoint: PubSubEndpoint = endpoint
80 self._uptime = datetime.utcnow()
81 self._workers_count = (lambda envar: int(envar) if envar.isdigit() else 1)(
82 os.environ.get("UVICORN_NUM_WORKERS", "1")
83 )
85 # helps us realize when another server already responded to a sync request
86 self._worker_id = uuid4().hex
88 # state: Dict[str, List[ChannelStats]]
89 # The state is built in this way so it will be easy to understand how much OPAL clients (vs. rpc clients)
90 # you have connected to your OPAL server and to help merge client lists between servers.
91 # The state is keyed by unique client id (A unique id that each opal client can set in env var `OPAL_CLIENT_STAT_ID`)
92 self._state: ServerStats = ServerStats(
93 uptime=self._uptime,
94 clients={},
95 servers={self._worker_id},
96 version=module_version(opal_server.__name__),
97 )
99 # rpc_id_to_client_id:
100 # dict to help us get client id without another loop
101 self._rpc_id_to_client_id: Dict[str, str] = {}
102 self._lock = asyncio.Lock()
103 self._synced_after_wakeup = asyncio.Event()
104 self._received_sync_messages: Set[str] = set()
105 self._publish_tasks = TasksPool()
106 self._seen_servers: Dict[str, datetime] = {}
107 self._periodic_keepalive_task: asyncio.Task | None = None
109 @property
110 def state(self) -> ServerStats:
111 return self._state
113 @property
114 def state_brief(self) -> ServerStatsBrief:
115 return ServerStatsBrief(
116 uptime=self._state.uptime,
117 version=self._state.version,
118 client_count=len(self._state.clients),
119 server_count=len(self._state.servers) / self._workers_count,
120 )
122 async def _expire_old_servers(self):
123 async with self._lock:
124 now = datetime.utcnow()
125 still_alive = {}
126 for server_id, last_seen in self._seen_servers.items():
127 if (now - last_seen).total_seconds() < float( 127 ↛ 126line 127 didn't jump to line 126 because the condition on line 127 was always true
128 opal_server_config.STATISTICS_SERVER_KEEPALIVE_TIMEOUT
129 ):
130 still_alive[server_id] = last_seen
131 self._seen_servers = still_alive
132 self._state.servers = {self._worker_id} | set(self._seen_servers.keys())
134 async def _periodic_server_keepalive(self):
135 while True:
136 try:
137 await self._expire_old_servers()
138 self._publish(
139 opal_server_config.STATISTICS_SERVER_KEEPALIVE_CHANNEL,
140 ServerKeepalive(worker_id=self._worker_id).dict(),
141 )
142 await asyncio.sleep(
143 float(opal_server_config.STATISTICS_SERVER_KEEPALIVE_TIMEOUT) / 2
144 )
145 except asyncio.CancelledError:
146 logger.debug("Statistics: periodic server keepalive cancelled")
147 return
148 except Exception as e:
149 logger.exception("Statistics: periodic server keepalive failed")
150 logger.exception("Statistics: periodic server keepalive failed")
152 def _publish(self, channel: str, message: Any):
153 self._publish_tasks.add_task(self._endpoint.publish([channel], message))
155 async def run(self):
156 """Subscribe to two channels to be able to sync add and delete of
157 clients."""
158 await self._endpoint.subscribe(
159 [opal_server_config.STATISTICS_WAKEUP_CHANNEL],
160 self._receive_other_worker_wakeup_message,
161 )
162 await self._endpoint.subscribe(
163 [opal_server_config.STATISTICS_STATE_SYNC_CHANNEL],
164 self._receive_other_worker_synced_state,
165 )
166 await self._endpoint.subscribe(
167 [opal_server_config.STATISTICS_SERVER_KEEPALIVE_CHANNEL],
168 self._receive_other_worker_keepalive_message,
169 )
170 await self._endpoint.subscribe(
171 [opal_common_config.STATISTICS_ADD_CLIENT_CHANNEL], self._add_client
172 )
173 await self._endpoint.subscribe(
174 [opal_common_config.STATISTICS_REMOVE_CLIENT_CHANNEL],
175 self._sync_remove_client,
176 )
178 # wait before publishing the wakeup message, due to the fact we are
179 # counting on the broadcaster to listen and to replicate the message
180 # to the other workers / server nodes in the networks.
181 # However, since broadcaster is using asyncio.create_task(), there is a
182 # race condition that is mitigated by this asyncio.sleep() call.
183 await asyncio.sleep(SLEEP_TIME_FOR_BROADCASTER_READER_TO_START)
184 # Let all the other opal servers know that new opal server started
185 logger.info(f"sending stats wakeup message: {self._worker_id}")
186 self._publish(
187 opal_server_config.STATISTICS_WAKEUP_CHANNEL,
188 SyncRequest(requesting_worker_id=self._worker_id).dict(),
189 )
190 self._periodic_keepalive_task = asyncio.create_task(
191 self._periodic_server_keepalive()
192 )
194 async def stop(self):
195 if self._periodic_keepalive_task: 195 ↛ exitline 195 didn't return from function 'stop' because the condition on line 195 was always true
196 self._periodic_keepalive_task.cancel()
197 await self._periodic_keepalive_task
198 self._periodic_keepalive_task = None
200 async def _sync_remove_client(self, subscription: Subscription, rpc_id: str):
201 """Helper function to recall remove client in all servers.
203 Args:
204 subscription (Subscription): not used, we get it from callbacks.
205 rpc_id (str): channel id of rpc channel used as identifier to client id
206 """
208 await self.remove_client(rpc_id=rpc_id, topics=[], publish=False)
210 async def _receive_other_worker_wakeup_message(
211 self, subscription: Subscription, sync_request: dict
212 ):
213 """Callback when new server wakes up and requests our statistics state.
215 Sends state only if we have state of our own and another
216 response to that request was not already received. Always reply
217 with hello message to refresh the "workers" state of other
218 servers.
219 """
220 try:
221 request = SyncRequest(**sync_request)
222 except pydantic.ValidationError as e:
223 logger.warning(
224 f"Got invalid statistics sync request from another server, error: {repr(e)}"
225 )
226 return
228 if self._worker_id == request.requesting_worker_id: 228 ↛ 235line 228 didn't jump to line 235 because the condition on line 228 was always true
229 # skip my own requests
230 logger.debug(
231 f"IGNORING my own stats wakeup message: {request.requesting_worker_id}"
232 )
233 return
235 logger.debug(f"received stats wakeup message: {request.requesting_worker_id}")
237 if len(self._state.clients):
238 # wait random time in order to reduce the number of messages sent by all the other opal servers
239 await asyncio.sleep(uniform(MIN_TIME_TO_WAIT, MAX_TIME_TO_WAIT))
240 # if didn't get any other message it means that this server is the first one to pass the sleep
241 if request.requesting_worker_id not in self._received_sync_messages:
242 logger.info(
243 f"[{request.requesting_worker_id}] respond with my own stats"
244 )
245 self._publish(
246 opal_server_config.STATISTICS_STATE_SYNC_CHANNEL,
247 SyncResponse(
248 requesting_worker_id=request.requesting_worker_id,
249 clients=self._state.clients,
250 rpc_id_to_client_id=self._rpc_id_to_client_id,
251 ).dict(),
252 )
254 async def _receive_other_worker_synced_state(
255 self, subscription: Subscription, sync_response: dict
256 ):
257 """Callback when another server sends us it's statistics data as a
258 response to a sync request.
260 Args:
261 subscription (Subscription): not used, we get it from callbacks.
262 rpc_id (Dict[str, List[ChannelStats]]): state from remote server
263 """
264 try:
265 response = SyncResponse(**sync_response)
266 except pydantic.ValidationError as e:
267 logger.warning(
268 f"Got invalid statistics sync response from another server, error: {repr(e)}"
269 )
270 return
272 async with self._lock:
273 self._received_sync_messages.add(response.requesting_worker_id)
275 # update my state only if this server don't have a state
276 if not len(self._state.clients) and not self._synced_after_wakeup.is_set():
277 logger.info(f"[{response.requesting_worker_id}] applying server stats")
278 self._state.clients = response.clients
279 self._rpc_id_to_client_id = response.rpc_id_to_client_id
280 self._synced_after_wakeup.set()
282 async def _receive_other_worker_keepalive_message(
283 self, subscription: Subscription, keepalive_message: dict
284 ):
285 async with self._lock:
286 self._seen_servers[keepalive_message["worker_id"]] = datetime.now()
287 self._state.servers.add(keepalive_message["worker_id"])
289 async def _add_client(self, subscription: Subscription, stats_message: dict):
290 """Add client record to statistics state.
292 Args:
293 subscription (Subscription): not used, we get it from callbacks.
294 stat_msg (ChannelStats): statistics data for channel, rpc_id - channel identifier; client_id - client identifier
295 """
296 try:
297 stats = ChannelStats(**stats_message)
298 except pydantic.ValidationError as e:
299 logger.warning(
300 f"Got invalid statistics message from client, error: {repr(e)}"
301 )
302 return
303 try:
304 client_id = stats.client_id
305 rpc_id = stats.rpc_id
306 logger.info(
307 "Set client statistics {client_id} on channel {rpc_id} with {topics}",
308 client_id=client_id,
309 rpc_id=rpc_id,
310 topics=", ".join(stats.topics),
311 )
312 async with self._lock:
313 self._rpc_id_to_client_id[rpc_id] = client_id
314 if client_id in self._state.clients:
315 # Limiting the number of channels per client to avoid memory issues if client opens too many channels
316 if (
317 len(self._state.clients[client_id])
318 < opal_server_config.MAX_CHANNELS_PER_CLIENT
319 ):
320 self._state.clients[client_id].append(stats)
321 else:
322 logger.warning(
323 f"Client '{client_id}' reached the maximum number of open RPC channels"
324 )
325 else:
326 self._state.clients[client_id] = [stats]
327 except Exception as err:
328 logger.exception("Add client to server statistics failed")
330 async def remove_client(self, rpc_id: str, topics: TopicList, publish=True):
331 """Remove client record from statistics state.
333 Args:
334 rpc_id (str): channel id of rpc channel used as identifier to client id
335 topics (TopicList): not used, we get it from callbacks.
336 publish (bool): used to stop republish cycle
337 """
338 if rpc_id not in self._rpc_id_to_client_id: 338 ↛ 344line 338 didn't jump to line 344 because the condition on line 338 was always true
339 logger.debug(
340 f"Statistics.remove_client() got unknown rpc id: {rpc_id} (probably broadcaster)"
341 )
342 return
344 try:
345 logger.info("Trying to remove {rpc_id} from statistics", rpc_id=rpc_id)
346 client_id = self._rpc_id_to_client_id[rpc_id]
347 for index, stats in enumerate(self._state.clients[client_id]):
348 if stats.rpc_id == rpc_id:
349 async with self._lock:
350 # remove the stats record matching the removed rpc id
351 del self._state.clients[client_id][index]
352 # remove the connection between rpc and client, once we removed it from state
353 del self._rpc_id_to_client_id[rpc_id]
354 # if no client records left in state remove the client entry
355 if not len(self._state.clients[client_id]):
356 del self._state.clients[client_id]
357 break
358 except Exception as err:
359 logger.warning(f"Remove client from server statistics failed: {repr(err)}")
360 # publish removed client so each server worker and server instance would get it
361 if publish:
362 logger.info(
363 "Publish rpc_id={rpc_id} to be removed from statistics",
364 rpc_id=rpc_id,
365 )
366 self._publish(
367 opal_common_config.STATISTICS_REMOVE_CLIENT_CHANNEL,
368 rpc_id,
369 )
372def init_statistics_router(stats: Optional[OpalStatistics] = None):
373 """Initializes a route where a client (or any other network peer) can
374 inquire what opal clients are currently connected to the server and on what
375 topics are they registered.
377 If the OPAL server does not have statistics enabled, the route will
378 return 501 Not Implemented
379 """
380 router = APIRouter()
382 @router.get("/statistics", response_model=ServerStats)
383 async def get_statistics():
384 """Route to serve server statistics."""
385 if stats is None: 385 ↛ 386line 385 didn't jump to line 386 because the condition on line 385 was never true
386 raise HTTPException(
387 status_code=status.HTTP_501_NOT_IMPLEMENTED,
388 detail={
389 "error": "This OPAL server does not have statistics turned on."
390 + " To turn on, set this config var: OPAL_STATISTICS_ENABLED=true"
391 },
392 )
393 logger.info("Serving statistics")
394 return stats.state
396 @router.get("/stats", response_model=ServerStatsBrief)
397 async def get_stat_counts():
398 """Route to serve only server and client instanace counts."""
399 if stats is None: 399 ↛ 400line 399 didn't jump to line 400 because the condition on line 399 was never true
400 raise HTTPException(
401 status_code=status.HTTP_501_NOT_IMPLEMENTED,
402 detail={
403 "error": "This OPAL server does not have statistics turned on."
404 + " To turn on, set this config var: OPAL_STATISTICS_ENABLED=true"
405 },
406 )
407 logger.info("Serving brief statistics info")
408 return stats.state_brief
410 return router