openrvdas.logger.writers.websocket_writer

  1#!/usr/bin/env python3
  2"""
  3"""
  4import asyncio
  5import logging
  6import ssl
  7import threading
  8import time
  9import inspect
 10
 11from typing import Union
 12from urllib.parse import urlparse
 13
 14
 15from logger.writers.writer import Writer  # noqa E402
 16
 17try:
 18    import websockets
 19
 20    WEBSOCKETS_INSTALLED = True
 21except ImportError:
 22    WEBSOCKETS_INSTALLED = False
 23
 24
 25################################################################################
 26class WebsocketWriter(Writer):
 27    ############################
 28
 29    def __init__(self, uri, cert_file=None, key_file=None, max_queue_size=None, **kwargs):
 30        """
 31        ```
 32        uri             Protocol, hostname and port to serve as. E.g. 'wss://openrvdas:8081'
 33                        If protocol is 'wss', use SSL, and cert_file and key_file must be
 34                        specified.
 35
 36        cert_file       If using ssl, the file path to relevant certificate and key files.
 37        key_file
 38
 39        max_queue_size  If specified, the maximum number of records to hold in each client's
 40                        send queue. When the queue exceeds this size, the oldest records are
 41                        discarded and a warning is logged indicating how many records were
 42                        dropped and how old the oldest dropped record was. Useful when
 43                        consumers only care about the most recent data (e.g. live dashboards)
 44                        and you want to avoid unbounded queue growth during network hiccups.
 45        ```
 46        """
 47        # Initialize type checking
 48        super().__init__(**kwargs)  # processes 'quiet' and type hints
 49
 50        if not WEBSOCKETS_INSTALLED:
 51            raise ImportError('WebsocketWriter requires Python "websockets" module; '
 52                              'please run "pip install websockets"')
 53        self.uri = uri
 54        parsed_uri = urlparse(uri)
 55        self.host = parsed_uri.hostname
 56        self.port = parsed_uri.port
 57        self.protocol = parsed_uri.scheme
 58
 59        if self.protocol == 'wss':
 60            self.ssl = True
 61            if (not cert_file or not key_file):
 62                raise ValueError('Both cert_file and key_file must be specified for wss')
 63        elif self.protocol == 'ws':
 64            self.ssl = False
 65            if (cert_file or key_file):
 66                raise ValueError('If protocol is ws, cert_file and key_file should be empty')
 67        else:
 68            raise ValueError(f'Protocol "{self.protocol}" not recognized. Must be ws or wss')
 69
 70        self.cert_file = cert_file
 71        self.key_file = key_file
 72        self.max_queue_size = max_queue_size
 73
 74        # Map from client_id to websocket
 75        self.client_map = {}
 76        self.client_map_lock = threading.Lock()
 77
 78        # Send queue for each client
 79        self.send_queue = {}
 80
 81        # Event loop - we'll set this in run()
 82        self.loop = None
 83
 84        # Start async websocket server in a separate thread
 85        self.server_run_thread = threading.Thread(target=self.run)
 86        self.server_run_thread.start()
 87
 88    ############################
 89    def _put_and_trim(self, client_id, record):
 90        """Enqueue a record (with timestamp) and drop oldest items if over max_queue_size.
 91
 92        Runs inside the event loop (called via call_soon_threadsafe), so asyncio
 93        Queue operations are safe here.
 94        """
 95        queue = self.send_queue.get(client_id)
 96        if queue is None:
 97            return  # client disconnected before this callback ran
 98
 99        queue.put_nowait((time.time(), record))
100
101        if self.max_queue_size and queue.qsize() > self.max_queue_size:
102            n_to_drop = queue.qsize() - self.max_queue_size
103            oldest_time = None
104            for _ in range(n_to_drop):
105                try:
106                    ts, _ = queue.get_nowait()
107                    if oldest_time is None:
108                        oldest_time = ts
109                except asyncio.QueueEmpty:
110                    break
111            age_str = (f', oldest was {time.time() - oldest_time:.1f}s old'
112                       if oldest_time is not None else '')
113            logging.warning(
114                f'WebsocketWriter: dropped {n_to_drop} queued record(s) for client '
115                f'{client_id}{age_str} (max_queue_size={self.max_queue_size})')
116
117    ############################
118    async def _send_from_queue(self, client_id):
119        """
120        Asynchronously wait for stuff to show up in client's queue, pop it off
121        and send to client's websocket.  Queue items are (enqueue_time, record) tuples.
122        """
123        try:
124            # The websocket, send_queue and send_queue_lock for this client
125            websocket = self.client_map[client_id]
126            send_queue = self.send_queue[client_id]
127            while True:
128                _enqueue_time, record = await send_queue.get()
129                await websocket.send(record)
130                logging.debug(f'WebsocketWriter sent client {client_id} record: {record}')
131
132        except websockets.exceptions.ConnectionClosed:  # type: ignore
133            logging.info(f'Websocket connection lost for client {client_id}')
134
135    ############################
136    # FIX: path=None makes this compatible with both new (1 arg) and old (2 args) libs
137    async def _websocket_handler(self, websocket, path=None):
138        with self.client_map_lock:
139            # Find a unique client_id for this client
140            client_id = 0
141            while client_id in self.client_map:
142                client_id += 1
143
144            logging.info(f'New client #{client_id} attached')
145            self.client_map[client_id] = websocket
146            self.send_queue[client_id] = asyncio.Queue()
147
148        # Task that produces data that we're going to send out on websocket
149        tasks = [asyncio.ensure_future(self._send_from_queue(client_id))]
150        await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
151
152        # When client disconnects, delete the websocket and queue it was using
153        with self.client_map_lock:
154            logging.info(f'WebsocketServer client #{client_id} completed')
155            del self.client_map[client_id]
156            del self.send_queue[client_id]
157
158    ############################
159    def run(self):
160        # Create an SSL context if we're using SSL
161        ssl_context = None
162        if self.ssl:
163            ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
164            ssl_context.load_cert_chain(
165                certfile=self.cert_file, keyfile=self.key_file)  # type: ignore
166
167        # Set up an event loop for the websocket server
168        self.loop = asyncio.new_event_loop()
169        asyncio.set_event_loop(self.loop)
170
171        # Build arguments dynamically to support both old and new websockets versions
172        serve_kwargs = {
173            'host': self.host,
174            'port': self.port,
175            'ssl': ssl_context
176        }
177
178        # Check signature of websockets.serve to decide on 'handler' vs 'ws_handler'
179        sig = inspect.signature(websockets.serve)
180        if 'handler' in sig.parameters:
181            serve_kwargs['handler'] = self._websocket_handler
182        else:
183            serve_kwargs['ws_handler'] = self._websocket_handler
184
185        async def runner():
186            # Create the server object/context manager
187            # We must do this INSIDE the coroutine so that 'websockets'
188            # can find the running loop.
189            server_result = websockets.serve(**serve_kwargs)
190
191            # Modern websockets (v10+): returns an AsyncContextManager
192            if hasattr(server_result, '__aenter__'):
193                async with server_result:
194                    # Keep the loop running while the server is active
195                    await asyncio.Future()
196
197            # Legacy websockets: returns an awaitable that yields the server
198            else:
199                await server_result
200                # Keep the loop running
201                await asyncio.Future()
202
203        try:
204            # run_until_complete starts the loop, then executes runner(),
205            # which keeps running via await asyncio.Future()
206            self.loop.run_until_complete(runner())
207        except Exception as e:
208            logging.error(f"WebsocketWriter server loop error: {e}")
209
210    ############################
211    def write(self, record: Union[str, bytes]):
212        """Write a record to all connected clients"""
213
214        # See if it's something we can process, and if not, try digesting
215        if not self.can_process_record(record):  # inherited from BaseModule()
216            self.digest_record(record)  # inherited from BaseModule()
217            return
218
219        logging.debug(f'WebsocketWriter received record: {record}')
220        with self.client_map_lock:
221            for client_id in self.client_map:
222                logging.debug(f'Pushing record to client {client_id}')
223                self.loop.call_soon_threadsafe(  # type: ignore
224                    self._put_and_trim, client_id, record)