openrvdas.logger.writers.cached_data_writer

  1#!/usr/bin/env python3
  2
  3import asyncio
  4import json
  5import logging
  6import ssl
  7import threading
  8
  9from typing import Union
 10try:
 11    import websockets
 12    WEBSOCKETS_INSTALLED = True
 13except ImportError:
 14    WEBSOCKETS_INSTALLED = False
 15
 16from logger.writers.writer import Writer  # noqa: E402
 17from logger.utils.das_record import DASRecord  # noqa: E402
 18
 19
 20class CachedDataWriter(Writer):
 21    def __init__(self, data_server, start_server=False, back_seconds=480,
 22                 cleanup_interval=6, update_interval=1,
 23                 max_backup=60 * 60 * 24, use_wss=False, check_cert=False, **kwargs):
 24        """Feed passed records to a CachedDataServer via a websocket. Expects
 25        records in DASRecord or dict formats.
 26        ```
 27        data_server    [host:]port on which to look for data server
 28
 29        back_seconds   Number of seconds of back data to hold in cache
 30
 31        cleanup_interval   Remove old data every N seconds
 32
 33        update_interval    Serve updates to websocket clients every N seconds
 34
 35        max_backup    If the writer isn't able to connect to the data server,
 36                      it will locally cache records until it can. To avoid
 37                      unbounded memory usage, if max_backup is nonzero, it will
 38                      cache at most max_backup records before dropping the
 39                      oldest records. By default, cache one day's worth of
 40                      records at 1 Hz (86,400 records). If max_backup is zero,
 41                      cache size is unbounded.
 42
 43        use_wss -     If True, use secure websockets
 44
 45        check_cert  - If True and use_wss is True, check the server's TLS certificate
 46                      for validity; if a str, use as local filepath location of .pem
 47                      file to check against.
 48
 49        ```
 50        """
 51        if not WEBSOCKETS_INSTALLED:
 52            raise ImportError('CachedDataWriter requires Python "websockets" module; '
 53                              'please run "pip install websockets"')
 54
 55        super().__init__(**kwargs)  # processes 'quiet' and type hints
 56
 57        host_port = data_server.split(':')
 58        if len(host_port) == 1:
 59            self.data_server = 'localhost:' + data_server  # they gave us '8766'
 60        elif not len(host_port[0]):
 61            self.data_server = 'localhost' + data_server   # they gave us ':8766'
 62        else:
 63            self.data_server = data_server                 # they gave us 'host:8766'
 64
 65        self.websocket = None
 66        self.back_seconds = back_seconds
 67        self.cleanup_interval = cleanup_interval
 68        self.use_wss = use_wss
 69        self.check_cert = check_cert
 70        self.event_loop = asyncio.new_event_loop()
 71
 72        # "loop" parameter removed in Ubuntu 22, but needed in earlier releases
 73        try:
 74            self.send_queue = asyncio.Queue(maxsize=max_backup)
 75        except RuntimeError:
 76            self.send_queue = asyncio.Queue(maxsize=max_backup, loop=self.event_loop)
 77
 78        # Start the thread that will asynchronously pull stuff from the
 79        # queue and send to the websocket. Also will, if we've got our oue
 80        # data server, run cleanup from time to time.
 81        self.cached_data_writer_thread = threading.Thread(
 82            name='cached_data_writer_thread',
 83            target=self._cached_data_writer_loop, daemon=True)
 84        self.cached_data_writer_thread.start()
 85
 86    ############################
 87    def _cached_data_writer_loop(self):
 88        """Use an inner async function to pull stuff from the queue and send
 89        to the websocket. Also, if we've got our oue data server, run
 90        cleanup from time to time.
 91        """
 92
 93        ############################
 94        async def _async_send_records_loop(self):
 95            """Inner async function that actually does websocket writes
 96            and cleanups.
 97            """
 98            while True:
 99                logging.debug('CachedDataWriter trying to connect to '
100                              + self.data_server)
101                try:
102                    if self.use_wss:
103                        # If check_cert is a str, take it as the location of the
104                        # .pem file we'll check for validity. Otherwise, if not
105                        # False, take as a bool to verify by own means.
106                        ws_data_server = 'wss://' + self.data_server
107                        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
108                        if self.check_cert:
109                            if isinstance(self.check_cert, str):
110                                ssl_context.load_verify_locations(self.check_cert)
111                            else:
112                                ssl_context.verify_mode = ssl.CERT_REQUIRED
113                        else:
114                            ssl_context.verify_mode = ssl.CERT_NONE
115
116                    else:  # not using wss
117                        ws_data_server = 'ws://' + self.data_server
118                        ssl_context = None
119
120                    logging.debug(f'CachedDataWriter connecting to {ws_data_server}')
121                    async with websockets.connect(ws_data_server, ssl=ssl_context) as ws:
122                        logging.debug(f'Connected to data server {ws_data_server}')
123                        while True:
124                            try:
125                                record = self.send_queue.get_nowait()
126                                logging.debug('sending record: %s', record)
127                                record = {'type': 'publish', 'data': record}
128                                await ws.send(json.dumps(record))
129                                response = await ws.recv()
130                                logging.debug('received response: %s', response)
131                            except asyncio.QueueEmpty:
132                                await asyncio.sleep(.2)
133
134                except BrokenPipeError:
135                    pass
136                except AttributeError as e:
137                    logging.warning('CachedDataWriter websocket loop error: %s', e)
138                    await asyncio.sleep(0.1)
139                except websockets.exceptions.ConnectionClosed:
140                    logging.warning('CachedDataWriter lost websocket connection to '
141                                    'data server; trying to reconnect.')
142                    await asyncio.sleep(0.2)
143
144                except websockets.exceptions.InvalidStatusCode:
145                    logging.warning('CachedDataWriter InvalidStatusCode connecting to '
146                                    'data server; trying to reconnect.')
147                    await asyncio.sleep(0.2)
148
149                # If the websocket connection failed
150                except OSError as e:
151                    logging.warning('CachedDataWriter websocket connection to %s '
152                                    'failed; sleeping before trying again: %s',
153                                    self.data_server, str(e))
154                    await asyncio.sleep(5)
155
156        # Now call the async process in its own event loop
157        self.event_loop.run_until_complete(_async_send_records_loop(self))
158        self.event_loop.close()
159
160    ############################
161    def write(self, record: Union[DASRecord, dict]):
162        """Write out record. Expects passed records to be in one of three
163        formats:
164
165        1) DASRecord
166
167        2) a list of DASRecords
168
169        3) a dict encoding optionally a source data_id and timestamp and a
170           mandatory 'fields' key of field_name: value pairs. This is the format
171           emitted by default by ParseTransform:
172
173           {
174             'data_id': ...,
175             'timestamp': ...,
176             'fields': {
177               field_name: value,    # use default timestamp of 'now'
178               field_name: value,
179               ...
180             }
181           }
182
183        A twist on format (3) THAT WE'RE PROBABLY GOING TO PHASE OUT IN
184        FAVOR OF (2) is that the values may either be a singleton (int,
185        float, string, etc) or a list. If the value is a singleton, it is
186        taken at face value. If it is a list, it is assumed to be a list
187        of (value, timestamp) tuples, in which case the top-level
188        timestamp, if any, is ignored.
189
190           {
191             'data_id': ...,
192             'timestamp': ...,
193             'fields': {
194                field_name: [(timestamp, value), (timestamp, value),...],
195                field_name: [(timestamp, value), (timestamp, value),...],
196                ...
197             }
198           }
199
200        """
201        # See if it's something we can process, and if not, try digesting
202        if not self.can_process_record(record):  # inherited from BaseModule()
203            self.digest_record(record)  # inherited from BaseModule()
204            return
205
206        # Convert to a dict - inefficient, I know...
207        if isinstance(record, DASRecord):
208            record = json.loads(record.as_json())
209        if isinstance(record, dict):
210            # If our local queue is full, throw away the oldest entries
211            while self.send_queue.full():
212                try:
213                    logging.debug('CachedDataWriter queue full - dropping oldest...')
214                    self.send_queue.get_nowait()
215                except asyncio.QueueEmpty:
216                    logging.warning('CachedDataWriter queue is both full and empty?!?')
217
218            # Enqueue our latest record for send
219            try:
220                self.send_queue.put_nowait(record)
221            except asyncio.queues.QueueFull:
222                logging.warning('CachedDataWriter unable to write: write queue full')
223        else:
224            if not self.quiet:
225                logging.warning('CachedDataWriter got non-dict/DASRecord object of '
226                                'type %s: %s', type(record), str(record))